Disable non blocking on mps.

comfyanonymous · Dec 10, 2023 · 340177e · 340177e
1 parent 614b7e7
commit 340177e
Showing 1 changed file with 8 additions and 4 deletions.
diff --git a/comfy/model_management.py b/comfy/model_management.py
@@ -553,15 +553,19 @@ def cast_to_device(tensor, device, dtype, copy=False):
  elif is_intel_xpu():
  device_supports_cast = True
 
+ non_blocking = True
+ if is_device_mps(device):
+ non_blocking = False #pytorch bug? mps doesn't support non blocking
+
  if device_supports_cast:
  if copy:
  if tensor.device == device:
- return tensor.to(dtype, copy=copy, non_blocking=True)
- return tensor.to(device, copy=copy, non_blocking=True).to(dtype, non_blocking=True)
+ return tensor.to(dtype, copy=copy, non_blocking=non_blocking)
+ return tensor.to(device, copy=copy, non_blocking=non_blocking).to(dtype, non_blocking=non_blocking)
  else:
- return tensor.to(device, non_blocking=True).to(dtype, non_blocking=True)
+ return tensor.to(device, non_blocking=non_blocking).to(dtype, non_blocking=non_blocking)
  else:
- return tensor.to(device, dtype, copy=copy, non_blocking=True)
+ return tensor.to(device, dtype, copy=copy, non_blocking=non_blocking)
 
 def xformers_enabled():
  global directml_enabled