Update modeling_stdit.py

modeling_stdit.py

@@ -109,6 +109,10 @@ class STDiT(PreTrainedModel):
         Returns:
             x (torch.Tensor): output latent representation; of shape [B, C, T, H, W]
         """
+        x = x.to(self.final_layer.linear.weight.dtype)
+        timestep = timestep.to(self.final_layer.linear.weight.dtype)
+        y = y.to(self.final_layer.linear.weight.dtype)
+        
         # embedding
         x = self.x_embedder(x)  # [B, N, C]
         x = rearrange(x, "B (T S) C -> B T S C", T=self.num_temporal, S=self.num_spatial)