sever

server.py

@@ -38,13 +38,13 @@ parser.add_argument('--init-mode', type = str, default='single', help = 'How to
 args = parser.parse_args()
 args = utils.set_hyperparams(args)
 
-# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-device = torch.device("cpu")    
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+# device = torch.device("cpu")    
 args.device = device
 args.device1 = device
-# if torch.cuda.device_count() >= 2:
-#     args.device = "cuda:0"
-#     args.device1 = "cuda:1"
+if torch.cuda.device_count() >= 2:
+    args.device = "cuda:0"
+    args.device1 = "cuda:1"
 
 utils.seed_all(args.seed)
 np.set_printoptions(precision=5)