From 1d2eca2a3bffb64403d80753096188593ddb6f47 Mon Sep 17 00:00:00 2001
From: Sultan Abu Ghazal <sultanmobeen@hotmail.com>
Date: Thu, 17 Feb 2022 02:09:49 +0400
Subject: [PATCH] Set the gpu_nums even if the device_count is 1 or less

Set the gpu_nums even if the device_count is 1 or less. This line has to be there when training on a single gpu.
---
 cityscapes_train.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/cityscapes_train.py b/cityscapes_train.py
index 19a4a22..702b2a7 100644
--- a/cityscapes_train.py
+++ b/cityscapes_train.py
@@ -175,6 +175,7 @@ def train_model(args):
             model = torch.nn.DataParallel(model).cuda()  #multi-card data parallel
         else:
             print("single GPU for training")
+            args.gpu_nums = torch.cuda.device_count()
             model = model.cuda()  #1-card data parallel
     
     args.savedir = ( args.savedir + args.dataset + '/'+ network_type +"_M"+ str(M) + 'N' +str(N) + 'bs'