Don't pin tensor if not a torch.nn.parameter.Parameter (#10718)

This commit is contained in:
comfyanonymous
2025-11-11 16:33:30 -08:00
committed by GitHub
parent 5ebcab3c7d
commit 1199411747

View File

@@ -1107,6 +1107,9 @@ def pin_memory(tensor):
if MAX_PINNED_MEMORY <= 0:
return False
if type(tensor) is not torch.nn.parameter.Parameter:
return False
if not is_device_cpu(tensor.device):
return False
@@ -1116,6 +1119,9 @@ def pin_memory(tensor):
#on the GPU async. So dont trust the CUDA API and guard here
return False
if not tensor.is_contiguous():
return False
size = tensor.numel() * tensor.element_size()
if (TOTAL_PINNED_MEMORY + size) > MAX_PINNED_MEMORY:
return False