Fix default device setting

flairNLP · Nov 30, 2024 · 7edf0e4 · 7edf0e4
1 parent 1b416b3
commit 7edf0e4
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/transformer_ranker/ranker.py b/transformer_ranker/ranker.py
@@ -72,6 +72,8 @@ def run(
         """
         self._confirm_ranker_setup(estimator=estimator, layer_aggregator=layer_aggregator)
 
+        device = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu"))
+
         # Load all transformers into hf cache
         self._preload_transformers(models, device)