Update ldm/modules/diffusionmodules/flag_large_dit.py

ldm/modules/diffusionmodules/flag_large_dit.py

@@ -241,8 +241,11 @@ class TxtFlagLargeDiT(nn.Module):
 
         print(f"theta {theta} rope scaling {rope_scaling_factor} ntk {ntk_factor}")
 
+        # freqs = 1.0 / (theta ** (
+        #     torch.arange(0, dim, 2)[: (dim // 2)].float().cuda() / dim
+        # ))
         freqs = 1.0 / (theta ** (
-            torch.arange(0, dim, 2)[: (dim // 2)].float().cuda() / dim
+            torch.arange(0, dim, 2)[: (dim // 2)].float() / dim
         ))
         t = torch.arange(end, device=freqs.device, dtype=torch.float)  # type: ignore
         t = t / rope_scaling_factor