load model using device.index

anthdr · Jan 17, 2025 · 0f3ea55 · 0f3ea55
1 parent 484b745
commit 0f3ea55
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/eole/bin/tools/LM_scoring.py b/eole/bin/tools/LM_scoring.py
@@ -61,7 +61,7 @@ def run(cls, args):
         if len(config.gpu_ranks) > 1:
             logger.warning(f"gpu_ranks is {str(config.gpu_ranks)} but only the first one will be used.")
 
-        vocabs, model, model_opt = BaseModel.load_test_model(config, 0)
+        vocabs, model, model_opt = BaseModel.load_test_model(config, device.index)
         pad_token = vocabs["specials"].get("pad_token", DefaultTokens.PAD)
         padding_idx = vocabs["tgt"].tokens_to_ids[pad_token]
         criterion = torch.nn.CrossEntropyLoss(ignore_index=padding_idx, reduction="none")