Add code for multi-gpu HuggingfaceSubject

brain-score · Jul 11, 2023 · 44d030c · 44d030c
1 parent 7875581
commit 44d030c
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/brainscore_language/model_helpers/huggingface.py b/brainscore_language/model_helpers/huggingface.py
@@ -42,9 +42,8 @@ def __init__(
         self._logger = logging.getLogger(fullname(self))
         self.model_id = model_id
         self.region_layer_mapping = region_layer_mapping
-        self.basemodel = (model if model is not None else AutoModelForCausalLM.from_pretrained(self.model_id))
+        self.basemodel = (model if model is not None else AutoModelForCausalLM.from_pretrained(self.model_id, device_map="auto"))
         self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
-        self.basemodel.to(self.device)
         self.tokenizer = tokenizer if tokenizer is not None else AutoTokenizer.from_pretrained(self.model_id,
                                                                                                truncation_side='left')
         self.current_tokens = None  # keep track of current tokens

diff --git a/pyproject.toml b/pyproject.toml
@@ -21,6 +21,7 @@ dependencies = [
     "transformers>=4.11.3",
     "gensim",
     "joblib",
+    "accelerate",
 ]
 
 [project.optional-dependencies]