From 340177e6e85d076ab9e222e4f3c6a22f1fb4031f Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Sun, 10 Dec 2023 01:30:35 -0500 Subject: [PATCH] Disable non blocking on mps. --- comfy/model_management.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/comfy/model_management.py b/comfy/model_management.py index ef9bec54..0c51eee5 100644 --- 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