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

Fix onnx export.

Browse Source
This commit is contained in:
comfyanonymous 2024-08-23 17:52:47 -04:00
parent 7df42b9a23
commit 8ae23d8e80
1 changed files with 7 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
comfy/ops.py
View File

@ -20,9 +20,13 @@ import torch
import comfy.model_management import comfy.model_management
from comfy.cli_args import args from comfy.cli_args import args
def cast_to(weight, dtype=None, device=None, non_blocking=False, copy=True): def cast_to(weight, dtype=None, device=None, non_blocking=False, copy=False):
if not copy and (dtype is None or weight.dtype == dtype) and (device is None or weight.device == device): if device is None or weight.device == device:
return weight if not copy:
if dtype is None or weight.dtype == dtype:
return weight
return weight.to(dtype=dtype, copy=copy)
r = torch.empty_like(weight, dtype=dtype, device=device) r = torch.empty_like(weight, dtype=dtype, device=device)
r.copy_(weight, non_blocking=non_blocking) r.copy_(weight, non_blocking=non_blocking)
return r return r