Model: Cast autosplit_reserve to int

Torch errors if float values are passed (because bytes are not float types). Therefore, overestimate and cast to an int type. Resolves #97 Signed-off-by: kingbri <[email protected]>
theroyallab · Apr 22, 2024 · 88b0b6f · 88b0b6f
1 parent cab789e
commit 88b0b6f
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/backends/exllamav2/model.py b/backends/exllamav2/model.py
@@ -1,6 +1,7 @@
 """The model container class for ExLlamaV2 models."""
 
 import gc
+import math
 import pathlib
 import threading
 import time
@@ -130,7 +131,10 @@ def progress(loaded_modules: int, total_modules: int,
 
             autosplit_reserve_megabytes = unwrap(kwargs.get("autosplit_reserve"), [96])
             self.autosplit_reserve = list(
-                map(lambda value: value * 1024**2, autosplit_reserve_megabytes)
+                map(
+                    lambda value: int(math.ceil(value * 1024**2)),
+                    autosplit_reserve_megabytes,
+                )
             )
         elif gpu_count > 1:
             # Manual GPU split