File tree Expand file tree Collapse file tree 2 files changed +2
-2
lines changed
brainscore_language/model_helpers Expand file tree Collapse file tree 2 files changed +2
-2
lines changed Original file line number Diff line number Diff line change @@ -42,9 +42,8 @@ def __init__(
42
42
self ._logger = logging .getLogger (fullname (self ))
43
43
self .model_id = model_id
44
44
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" ))
46
46
self .device = 'cuda' if torch .cuda .is_available () else 'cpu'
47
- self .basemodel .to (self .device )
48
47
self .tokenizer = tokenizer if tokenizer is not None else AutoTokenizer .from_pretrained (self .model_id ,
49
48
truncation_side = 'left' )
50
49
self .current_tokens = None # keep track of current tokens
Original file line number Diff line number Diff line change @@ -21,6 +21,7 @@ dependencies = [
21
21
" transformers>=4.11.3" ,
22
22
" gensim" ,
23
23
" joblib" ,
24
+ " accelerate" ,
24
25
]
25
26
26
27
[project .optional-dependencies ]
You can’t perform that action at this time.
0 commit comments