Update modeling_tiny_mixtral.py

modeling_tiny_mixtral.py

@@ -582,9 +582,21 @@ class TinyMixtralForCausalLM(PreTrainedModel, GenerationMixin):
         )
         self.model = tiny_mixtral(args=args)
         self.config = config
-        self.model.to(config.device)
+        self._move_model_to_device(config.device)
         self.post_init()
 
+    def _move_model_to_device_safe(self, target_device):
+        """Safely move model to target device"""
+        # Check if any parameter is on meta device
+        has_meta_params = any(param.device.type == 'meta' for param in self.model.parameters())
+        
+        if has_meta_params:
+            print("Detected meta tensors, using to_empty() to move model")
+            self.model.to_empty(device=target_device)
+        else:
+            print(f"Moving model to {target_device} using standard to()")
+            self.model.to(target_device)
+
     def forward(self, input_ids, attention_mask=None, labels=None, **kwargs):
 
         outputs, load_balancing_loss = self.model(input_ids, start_pos=0)