Skip to content

Commit 1199411

Browse files
Don't pin tensor if not a torch.nn.parameter.Parameter (#10718)
1 parent 5ebcab3 commit 1199411

File tree

1 file changed

+6
-0
lines changed

1 file changed

+6
-0
lines changed

comfy/model_management.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1107,6 +1107,9 @@ def pin_memory(tensor):
11071107
if MAX_PINNED_MEMORY <= 0:
11081108
return False
11091109

1110+
if type(tensor) is not torch.nn.parameter.Parameter:
1111+
return False
1112+
11101113
if not is_device_cpu(tensor.device):
11111114
return False
11121115

@@ -1116,6 +1119,9 @@ def pin_memory(tensor):
11161119
#on the GPU async. So dont trust the CUDA API and guard here
11171120
return False
11181121

1122+
if not tensor.is_contiguous():
1123+
return False
1124+
11191125
size = tensor.numel() * tensor.element_size()
11201126
if (TOTAL_PINNED_MEMORY + size) > MAX_PINNED_MEMORY:
11211127
return False

0 commit comments

Comments
 (0)