Move to device when not accelerating dataloader

flux_train.py

@@ -30,7 +30,7 @@
 
 init_ipex()
 
-from accelerate.utils import set_seed
+from accelerate.utils import set_seed, send_to_device
 from library import deepspeed_utils, flux_train_utils, flux_utils, strategy_base, strategy_flux
 from library.sd3_train_utils import FlowMatchEulerDiscreteScheduler
 
@@ -648,6 +648,10 @@ def grad_hook(parameter: torch.Tensor):
                         accelerator.print("NaN found in latents, replacing with zeros")
                         latents = torch.nan_to_num(latents, 0, out=latents)
 
+                if args.no_accelerate_dataloader:
+                    #The accelerator wrapper of the dataloader was handling device placement, but we disabled it.
+                    batch = send_to_device(batch, accelerator.device)
+
                 text_encoder_outputs_list = batch.get("text_encoder_outputs_list", None)
                 if text_encoder_outputs_list is not None:
                     text_encoder_conds = text_encoder_outputs_list