Small optimization.

comfyanonymous · Jun 15, 2024 · 1281f93 · 1281f93
1 parent f2e844e
commit 1281f93
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/comfy/ldm/modules/diffusionmodules/mmdit.py b/comfy/ldm/modules/diffusionmodules/mmdit.py
@@ -243,9 +243,9 @@ def timestep_embedding(t, dim, max_period=10000):
         half = dim // 2
         freqs = torch.exp(
             -math.log(max_period)
-            * torch.arange(start=0, end=half, dtype=torch.float32)
+            * torch.arange(start=0, end=half, dtype=torch.float32, device=t.device)
             / half
-        ).to(device=t.device)
+        )
         args = t[:, None].float() * freqs[None]
         embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
         if dim % 2: