Skip to content

Commit 44d030c

Browse files
committed
Add code for multi-gpu HuggingfaceSubject
1 parent 7875581 commit 44d030c

File tree

2 files changed

+2
-2
lines changed

2 files changed

+2
-2
lines changed

brainscore_language/model_helpers/huggingface.py

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -42,9 +42,8 @@ def __init__(
4242
self._logger = logging.getLogger(fullname(self))
4343
self.model_id = model_id
4444
self.region_layer_mapping = region_layer_mapping
45-
self.basemodel = (model if model is not None else AutoModelForCausalLM.from_pretrained(self.model_id))
45+
self.basemodel = (model if model is not None else AutoModelForCausalLM.from_pretrained(self.model_id, device_map="auto"))
4646
self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
47-
self.basemodel.to(self.device)
4847
self.tokenizer = tokenizer if tokenizer is not None else AutoTokenizer.from_pretrained(self.model_id,
4948
truncation_side='left')
5049
self.current_tokens = None # keep track of current tokens

pyproject.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ dependencies = [
2121
"transformers>=4.11.3",
2222
"gensim",
2323
"joblib",
24+
"accelerate",
2425
]
2526

2627
[project.optional-dependencies]

0 commit comments

Comments
 (0)