From 2be27427116c58b027c0d5e1a45be435910f7c77 Mon Sep 17 00:00:00 2001
From: comfyanonymous <comfyanonymous@protonmail.com>
Date: Thu, 17 Aug 2023 01:58:54 -0400
Subject: [PATCH] Fix issue with regular torch version.

---
 comfy/model_management.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/comfy/model_management.py b/comfy/model_management.py
index 3736b57a..4c2441f4 100644
--- a/comfy/model_management.py
+++ b/comfy/model_management.py
@@ -390,7 +390,11 @@ def unet_inital_load_device(parameters, dtype):
         return torch_dev
 
     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_cpu = get_free_memory(cpu_dev)