Update ldm/modules/diffusionmodules/flag_large_dit.py

ldm/modules/diffusionmodules/flag_large_dit.py

@@ -241,12 +241,12 @@ 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() / dim
+            torch.arange(0, dim, 2)[: (dim // 2)].float().cuda() / dim
         ))
+        # freqs = 1.0 / (theta ** (
+        #     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
         freqs = torch.outer(t, freqs).float()  # type: ignore