change cuz model parallel was not working..

modular_Sewy2.py

@@ -2042,6 +2042,8 @@ class SewyV2ForCausalLM(SewyV2PreTrainedModel):
 
         s_z = self.s_z * (self.s_z_init/self.s_z_scale)
 
+        s_z = s_z.to(logits.device)
+
         logits = logits * s_z.view(1, 1, -1)
 
         loss = None