fix self.device

axolotl-ai-cloud · Jul 19, 2024 · a9b5ca4 · a9b5ca4
1 parent 0bfa6f7
commit a9b5ca4
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/axolotl/utils/models.py b/src/axolotl/utils/models.py
@@ -976,7 +976,7 @@ def load_model(self) -> Tuple[PreTrainedModel, Optional[PeftConfig]]:
 
         self.apply_lora_patch(model)
 
-        if "cuda" in self.device:
+        if "cuda" in self.device.__str__():
             for _ in range(3):
                 gc.collect()
                 torch.cuda.empty_cache()