--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…
x
Reference in New Issue
Block a user