From 81388b14f7f1857426cc31edbe83e5667bfabffe Mon Sep 17 00:00:00 2001
From: Izaak Curry <98251797+ScientistIzaak@users.noreply.github.com>
Date: Wed, 1 Jan 2025 22:33:07 -0800
Subject: [PATCH] add device parameter to TransformerModel

---
 src/smolagents/models.py | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/src/smolagents/models.py b/src/smolagents/models.py
index 6fc8dbb..6ad0ce9 100644
--- a/src/smolagents/models.py
+++ b/src/smolagents/models.py
@@ -284,7 +284,7 @@ class HfApiModel(Model):
 class TransformersModel(Model):
     """This engine initializes a model and tokenizer from the given `model_id`."""
 
-    def __init__(self, model_id: Optional[str] = None):
+    def __init__(self, model_id: Optional[str] = None, device: Optional[str] = None):
         super().__init__()
         default_model_id = "HuggingFaceTB/SmolLM2-1.7B-Instruct"
         if model_id is None:
@@ -293,15 +293,18 @@ class TransformersModel(Model):
                 f"`model_id`not provided, using this default tokenizer for token counts: '{model_id}'"
             )
         self.model_id = model_id
+        if device is None:
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.device = device
         try:
             self.tokenizer = AutoTokenizer.from_pretrained(model_id)
-            self.model = AutoModelForCausalLM.from_pretrained(model_id)
+            self.model = AutoModelForCausalLM.from_pretrained(model_id).to(self.device)
         except Exception as e:
             logger.warning(
                 f"Failed to load tokenizer and model for {model_id=}: {e}. Loading default tokenizer and model instead from {model_id=}."
             )
             self.tokenizer = AutoTokenizer.from_pretrained(default_model_id)
-            self.model = AutoModelForCausalLM.from_pretrained(default_model_id)
+            self.model = AutoModelForCausalLM.from_pretrained(default_model_id).to(self.device)
 
     def make_stopping_criteria(self, stop_sequences: List[str]) -> StoppingCriteriaList:
         class StopOnStrings(StoppingCriteria):