Update train.py

flyingsheepbin · May 30, 2022 · 7ed6133 · 7ed6133
1 parent 222bda1
commit 7ed6133
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/b7ns_stage1/train.py b/b7ns_stage1/train.py
@@ -14,7 +14,7 @@
 from torch_ema import ExponentialMovingAverage
 
 #os.environ['CUDA_VISIBLE_DEVICE'] = '1'
-os.environ["CUDA_VISIBLE_DEVICES"] = "1"
+os.environ["CUDA_VISIBLE_DEVICES"] = "0"
 
 def copy_parameters_to_model(copy_of_model_parameters, model):
     for s_param, param in zip(copy_of_model_parameters, model.parameters()):