--disable-smart-memory now disables loading model directly to vram.
This commit is contained in:
parent
c9b562aed1
commit
e9469e732d
|
@ -397,6 +397,9 @@ def unet_inital_load_device(parameters, dtype):
|
||||||
return torch_dev
|
return torch_dev
|
||||||
|
|
||||||
cpu_dev = torch.device("cpu")
|
cpu_dev = torch.device("cpu")
|
||||||
|
if DISABLE_SMART_MEMORY:
|
||||||
|
return cpu_dev
|
||||||
|
|
||||||
dtype_size = 4
|
dtype_size = 4
|
||||||
if dtype == torch.float16 or dtype == torch.bfloat16:
|
if dtype == torch.float16 or dtype == torch.bfloat16:
|
||||||
dtype_size = 2
|
dtype_size = 2
|
||||||
|
|
Loading…
Reference in New Issue