This commit is contained in:
lvmin
2023-08-09 14:00:59 -07:00
parent 573dc98eff
commit 0b72417588
2 changed files with 2 additions and 2 deletions
+1 -1
View File
@@ -980,7 +980,7 @@ class UNetModel(nn.Module):
self.num_classes is not None
), "must specify y if and only if the model is class-conditional"
hs = []
t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False)
t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False).to(x)
emb = self.time_embed(t_emb)
if self.num_classes is not None: