enable multi gpu support

labdao · Jul 10, 2022 · 6bb55e9 · 6bb55e9
1 parent 9cf6793
commit 6bb55e9
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/multiligand_inference.py b/multiligand_inference.py
@@ -106,7 +106,7 @@ def get_default_args(args, cmdline_args):
     return args
 
 def load_rec_and_model(args):
-    device = torch.device("cuda:0" if torch.cuda.is_available() and args.device == 'cuda' else "cpu")
+    device = torch.device("cuda" if torch.cuda.is_available() and args.device == 'cuda' else "cpu")
     print(f"device = {device}")
     # sys.exit()
     checkpoint = torch.load(args.checkpoint, map_location=device)
@@ -272,4 +272,4 @@ def main(arglist = None):
     write_while_inferring(lig_loader, model, args)
 
 if __name__ == '__main__':
-    main()
+    main()