moving fps tensor to gpu

Signed-off-by: Zeeshan Patel <[email protected]>
zpx01 · Nov 8, 2024 · d576cf7 · d576cf7
1 parent 9ca52c9
commit d576cf7
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/nemo/collections/diffusion/models/dit/dit_model.py b/nemo/collections/diffusion/models/dit/dit_model.py
@@ -224,7 +224,7 @@ def forward(
                 * B,
                 dtype=torch.bfloat16,
             ),
-        ).view(-1)
+        ).view(-1).cuda()
         if self.pre_process:
             # transpose to match
             x_B_S_D = self.x_embedder(x)