Add gpu choose option

scyonggg · Nov 10, 2022 · 47ca621 · 47ca621
1 parent 848fdd5
commit 47ca621
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/train.py b/train.py
@@ -41,6 +41,8 @@
     parser.add_argument('--run_name', type=str, default='scy_exp3', help='wandb run name')
     args = parser.parse_args()
 
+    os.environ['CUDA_VISIBLE_DEVICES'] = args.gpus
+
     #checkpoint directory
     checkpoint_path = os.path.join(settings.CHECKPOINT_PATH, args.net, settings.TIME_NOW)
     if not os.path.exists(checkpoint_path):
@@ -54,7 +56,7 @@
 
     if args.wandb:
         import wandb
-        wandb.init(project='test', entity="dnn_22_2", name=args.run_name, settings=wandb.Settings(code_dir="."))
+        wandb.init(project='scy_test', entity="dnn_22_2", name=args.run_name, settings=wandb.Settings(code_dir="."))
         wandb.run.log_code(".")
 
     #get dataloader