added missing function

terrierteam · Jul 16, 2023 · 862b973 · 862b973
1 parent f3452cf
commit 862b973
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/pyterrier_dr/__init__.py b/pyterrier_dr/__init__.py
@@ -1,4 +1,4 @@
-from .util import SimFn
+from .util import SimFn, infer_device
 from .indexes import DocnoFile, NilIndex, NumpyIndex, RankedLists, FaissFlat, FaissHnsw, MemIndex, TorchIndex
 from .flex import FlexIndex
 from .biencoder import BiEncoder, BiQueryEncoder, BiDocEncoder, BiScorer

diff --git a/pyterrier_dr/util.py b/pyterrier_dr/util.py
@@ -13,3 +13,8 @@ def wrapped(*args, **kwargs):
 
     def __init__(self, *args, **kwargs):
         return super().__init__(*args, **kwargs)
+
+def infer_device(device):
+    if device is None:
+        device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    return torch.device(device)