use same device for model and input

mobiusml · Oct 30, 2024 · 464e9d9 · 464e9d9
1 parent eb3e2a5
commit 464e9d9
Showing 1 changed file with 0 additions and 1 deletion.
diff --git a/aana/deployments/hqq_text_generation_deployment.py b/aana/deployments/hqq_text_generation_deployment.py
@@ -118,7 +118,6 @@ async def apply_config(self, config: dict[str, Any]):
                 self.dtype = Dtype.BFLOAT16
 
         if config_obj.quantize_on_fly:
-            self.model_kwargs["device_map"] = "cpu"
             self.model = AutoModelForCausalLM.from_pretrained(
                 self.model_id, torch_dtype=self.dtype.to_torch(), **self.model_kwargs
             )