quantization: instantiate model on CPU (#163)

timothylimyl · Apr 19, 2023 · 64636f1 · 64636f1
1 parent 06302ea
commit 64636f1
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/quantize.py b/quantize.py
@@ -187,8 +187,9 @@ def main(
     else:
         raise RuntimeError(f"unknown/unsupported quantization mode {quantize}")
 
+    # we avoid loading the entire model on the GPU and do this block by block
     with EmptyInitOnDevice(
-        device=device,
+        device="cpu",
         dtype=dtype,
     ):
         print("Loading model ...", file=sys.stderr)