type hint

microsoft · Jan 22, 2025 · 20e0680 · 20e0680
1 parent 716f52d
commit 20e0680
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/mattergen/common/utils/globals.py b/mattergen/common/utils/globals.py
@@ -14,7 +14,7 @@
 
 
 @lru_cache
-def get_device():
+def get_device() -> torch.device:
     if torch.cuda.is_available():
         return torch.device("cuda")
     if torch.backends.mps.is_available():
@@ -23,7 +23,7 @@ def get_device():
 
 
 @lru_cache
-def get_pyg_device():
+def get_pyg_device() -> torch.device:
     """
     Some operations of pyg don't work on MPS, so fall back to CPU.
     """