From f94cf433313b842e36740723247c4c8d5dce0957 Mon Sep 17 00:00:00 2001 From: kijai <40791699+kijai@users.noreply.github.com> Date: Mon, 4 Nov 2024 15:06:20 +0200 Subject: [PATCH] Update nodes.py --- nodes.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/nodes.py b/nodes.py index 75b33c7..3c1e947 100644 --- a/nodes.py +++ b/nodes.py @@ -181,7 +181,7 @@ class DownloadAndLoadMochiModel: vae_sd = load_torch_file(vae_path) if is_accelerate_available: for key in vae_sd: - set_module_tensor_to_device(vae, key, dtype=torch.float32, device=device, value=vae_sd[key]) + set_module_tensor_to_device(vae, key, dtype=torch.bfloat16, device=offload_device, value=vae_sd[key]) else: vae.load_state_dict(vae_sd, strict=True) vae.eval().to(torch.bfloat16).to("cpu")