device map

Author: mayank31398

src/pipelines/hf.py

@@ -9,10 +9,11 @@ class HF_Pipeline(Pipeline):
     def __init__(self, args: Namespace, device: str = "cpu") -> None:
         super().__init__(args)
 
-        model_kwargs = {"device_map": "auto"}
+        model_kwargs = {}
 
         if args.dtype == torch.int8:
             model_kwargs["load_in_8bit"] = True
+            model_kwargs["device_map"] = "auto"
         else:
             model_kwargs["torch_dtype"] = args.dtype