qwen-image splited training

examples/qwen_image/model_training/train.py

@@ -111,6 +111,7 @@ class QwenImageTrainingModule(DiffusionTrainingModule):
     
     def forward(self, data, inputs=None):
         if inputs is None: inputs = self.forward_preprocess(data)
+        else: inputs = self.transfer_data_to_device(inputs, self.pipe.device)
         models = {name: getattr(self.pipe, name) for name in self.pipe.in_iteration_models}
         loss = self.pipe.training_loss(**models, **inputs)
         return loss