Auto Selection of Device

aed0528c · ADRIAN AYUSO MUNOZ · aa6f5904 · aed0528c
Commit aed0528c authored Jan 30, 2023 by ADRIAN AYUSO MUNOZ
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

deepSnapPred.py deepSnapPred.py +1 -1

No files found.
--- a/deepSnapPred.py
+++ b/deepSnapPred.py
@@ -20,7 +20,7 @@ import warnings
 warnings.filterwarnings("ignore")

 edges = [('phenotype', 'dis_dru_the', 'drug')]
-device = 'cuda'
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

 # ---------------------------
 # FUNCTIONS