make onnx compatible

- remove broadcast_to - remove returning a dictionary
harishanand95 · Aug 29, 2022 · 8dd4e82 · 8dd4e82
1 parent efa773a
commit 8dd4e82
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/src/diffusers/models/unet_2d_condition.py b/src/diffusers/models/unet_2d_condition.py
@@ -134,7 +134,8 @@ def forward(
             timesteps = timesteps[None].to(sample.device)
 
         # broadcast to batch dimension
-        timesteps = timesteps.broadcast_to(sample.shape[0])
+        #timesteps = timesteps.broadcast_to(sample.shape[0])
+        timesteps = timesteps * torch.ones(sample.shape[0])
 
         t_emb = self.time_proj(timesteps)
         emb = self.time_embedding(t_emb)
@@ -181,6 +182,6 @@ def forward(
         sample = self.conv_act(sample)
         sample = self.conv_out(sample)
 
-        output = {"sample": sample}
+        #output = {"sample": sample}
 
-        return output
+        return sample