298977887/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2025-01-25 15:55:18 +00:00
Code Issues Packages Projects Releases Wiki Activity

Less broken non blocking?

Browse Source
This commit is contained in:
comfyanonymous 2024-08-18 16:53:17 -04:00
parent 6730f3e1a3
commit 39f114c44b
1 changed files with 3 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
comfy/ops.py
View File

@ -21,8 +21,9 @@ import comfy.model_management
def cast_to(weight, dtype=None, device=None, non_blocking=False):
non_blocking = False
return weight.to(device=device, dtype=dtype, non_blocking=non_blocking)
r = torch.empty_like(weight, dtype=dtype, device=device)
r.copy_(weight, non_blocking=non_blocking)
return r
def cast_to_input(weight, input, non_blocking=False):
return cast_to(weight, input.dtype, input.device, non_blocking=non_blocking)