Commit 2be27427 authored by comfyanonymous's avatar comfyanonymous
Browse files

Fix issue with regular torch version.

parent 89a0767a
...@@ -390,7 +390,11 @@ def unet_inital_load_device(parameters, dtype): ...@@ -390,7 +390,11 @@ def unet_inital_load_device(parameters, dtype):
return torch_dev return torch_dev
cpu_dev = torch.device("cpu") cpu_dev = torch.device("cpu")
model_size = dtype.itemsize * parameters dtype_size = 4
if dtype == torch.float16 or dtype == torch.bfloat16:
dtype_size = 2
model_size = dtype_size * parameters
mem_dev = get_free_memory(torch_dev) mem_dev = get_free_memory(torch_dev)
mem_cpu = get_free_memory(cpu_dev) mem_cpu = get_free_memory(cpu_dev)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment