Update pretrained.py: Quick fix to be able to load pertained_models d…

…irectly to GPU. (#254) * Update pretrained.py quick fix to be able to load pertained_models directly to GPU. * Update CHANGELOG.md
allenai · Apr 23, 2021 · 659c71f · 659c71f
1 parent 2ba85b5
commit 659c71f
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 0 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -18,6 +18,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 
+- `pretrained.load_predictor()` now allows for loading model onto GPU.
 - `VqaMeasure` now calculates correctly in the distributed case.
 - `ConllCorefScores` now calculates correctly in the distributed case.
 - `SrlEvalScorer` raises an appropriate error if run in the distributed setting.

diff --git a/allennlp_models/pretrained.py b/allennlp_models/pretrained.py
@@ -61,6 +61,7 @@ def get_pretrained_models() -> Dict[str, ModelCard]:
 def load_predictor(
     model_id: str,
     pretrained_models: Dict[str, ModelCard] = None,
+    cuda_device: int = -1,
     overrides: Union[str, Dict[str, Any]] = None,
 ) -> Predictor:
     """
@@ -76,5 +77,6 @@ def load_predictor(
     return Predictor.from_path(
         model_card.model_usage.archive_file,
         predictor_name=model_card.registered_predictor_name,
+        cuda_device=cuda_device,
         overrides=overrides,
     )