modified: src/flux/sampling.py

src/flux/sampling.py

@@ -522,6 +522,24 @@ def denoise_multi_turn_consistent(
         info['second_order'] = False
         info['inject'] = inject_list[i]
 
+     # Check and print devices of all input tensors
+    tensor_info = {
+        'img': img,
+        'img_ids': img_ids,
+        'txt': txt,
+        'txt_ids': txt_ids,
+        'y (vec)': vec, # Renamed for clarity in print
+        'timesteps (t_vec)': t_vec, # Renamed for clarity in print
+        'guidance (guidance_vec)': guidance_vec # Renamed for clarity in print
+    }
+    for name, tensor in tensor_info.items():
+        if torch.is_tensor(tensor):
+            print(f"{name} device: {tensor.device}, shape: {tensor.shape}")
+        else:
+            print(f"{name} is not a tensor: {type(tensor)} (value: {tensor})")
+
+
+
         if next_step_velocity is None:
             pred, info = model(
                 img=img,