update stdit dtype (#392)

hpcaitech · May 11, 2024 · ea41df3 · ea41df3
1 parent 09e53db
commit ea41df3
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/opensora/models/stdit/stdit.py b/opensora/models/stdit/stdit.py
@@ -267,10 +267,10 @@ def forward(self, x, timestep, y, mask=None, x_mask=None):
         Returns:
             x (torch.Tensor): output latent representation; of shape [B, C, T, H, W]
         """
-
-        x = x.to(self.dtype)
-        timestep = timestep.to(self.dtype)
-        y = y.to(self.dtype)
+        dtype = self.x_embedder.proj.weight.dtype
+        x = x.to(dtype)
+        timestep = timestep.to(dtype)
+        y = y.to(dtype)
 
         # embedding
         x = self.x_embedder(x)  # [B, N, C]