From 86a4c702de81cda7658d40ced761fe572d438c8b Mon Sep 17 00:00:00 2001 From: Kartik Pradeepan Date: Wed, 3 Jun 2026 04:53:20 -0400 Subject: [PATCH] skip localizer model.to(device) when model is device_map-sharded --- brainscore_language/model_helpers/localize.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/brainscore_language/model_helpers/localize.py b/brainscore_language/model_helpers/localize.py index 95b09897..a7cd4b7b 100644 --- a/brainscore_language/model_helpers/localize.py +++ b/brainscore_language/model_helpers/localize.py @@ -106,7 +106,8 @@ def extract_representations( logger.debug(f"> Using Device: {device}") model.eval() - model.to(device) + if not getattr(model, "hf_device_map", None): # skip when already sharded via device_map + model.to(device) final_layer_representations = { "sentences": {layer_name: np.zeros((len(langloc_dataset.sentences), hidden_dim)) for layer_name in layer_names},