Update main.py

ZhuLvs · Aug 28, 2024 · bbdd0a7 · bbdd0a7
1 parent 25639a4
commit bbdd0a7
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/model/main.py b/model/main.py
@@ -9,7 +9,7 @@
 from Inference import train_one_epoch, evaluate
 
 if __name__ == '__main__':
-    device = torch.device("cuda:4")
+    device = torch.device("cuda:0")
     model = CustomDeepLabV3().to(device)
     criterion = MaskedL1Loss()
     optimizer = torch.optim.Adam([