diff --git a/toolkit/memory_management/manager_modules.py b/toolkit/memory_management/manager_modules.py index 7dac4b59..1b158e3e 100644 --- a/toolkit/memory_management/manager_modules.py +++ b/toolkit/memory_management/manager_modules.py @@ -99,7 +99,10 @@ def _ensure_cpu_pinned(t: Optional[torch.Tensor]) -> Optional[torch.Tensor]: if t is None: return None if t.device.type != "cpu": - t = t.to("cpu", copy=True) + try: + t = t.to("cpu", copy=True) + except Exception: + t = t.to("cpu") # Don't attempt to pin quantized tensors; many backends don't support it if _is_quantized_tensor(t): return t