From 9f08b3cf3a38cf38f89401f67239a469246a4492 Mon Sep 17 00:00:00 2001
From: Artem Bolgar <artem@sesameai.com>
Date: Thu, 18 Apr 2024 20:43:05 -0700
Subject: [PATCH] Fixing quantize in int4 mode

---
 quantize.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/quantize.py b/quantize.py
index af17a698..0a91938b 100644
--- a/quantize.py
+++ b/quantize.py
@@ -539,7 +539,7 @@ def quantize(
     device: str = default_device,
 ) -> None:
     assert checkpoint_path.is_file(), checkpoint_path
-    device = 'cpu'
+    print(f"Using device={device}")
     precision = torch.bfloat16
 
     print("Loading model ...")