update

src/attribution/attntrace.py

@@ -41,7 +41,7 @@ class AttnTraceAttribution(Attribution):
         if self.llm.model!=None:
             self.model = self.llm.model
         else:
-            self.model = self.llm._load_model_if_needed()
+            self.model = self.llm._load_model_if_needed().to("cuda")
         self.layers = range(len(self.model.model.layers))
         model = self.model
         tokenizer = self.tokenizer

src/models/Llama.py

@@ -24,17 +24,18 @@ class Llama(Model):
         ]
 
     def _load_model_if_needed(self):
-        if self.model is None:
-            self.model = AutoModelForCausalLM.from_pretrained(
+        if self._model is None:
+            model = AutoModelForCausalLM.from_pretrained(
                 self.name,
                 torch_dtype=torch.bfloat16,
-                device_map=self.device,
-                token=self.hf_token
+                token=self.hf_token,
+                device_map="auto",  # or omit entirely to default to CPU
             )
-        return self.model
+            self._model = model
+        return self._model
 
     def query(self, msg, max_tokens=128000):
-        model = self._load_model_if_needed()
+        model = self._load_model_if_needed().to("cuda")
         messages = self.messages
         messages[1]["content"] = msg