From 838c7978de6f85ce7fe1affadbb9b6b60b4a633c Mon Sep 17 00:00:00 2001
From: Steve Nyemba <nyemba@gmail.com>
Date: Tue, 12 Apr 2022 14:32:39 -0500
Subject: [PATCH] bug fix: gpu visibility

---
 data/gan.py | 21 ++++++++++++---------
 1 file changed, 12 insertions(+), 9 deletions(-)

diff --git a/data/gan.py b/data/gan.py
index f5705ea..e0f97b1 100644
--- a/data/gan.py
+++ b/data/gan.py
@@ -61,16 +61,19 @@ class GNet :
                 self.logs = {}
 
                 # self.NUM_GPUS = 1 if 'num_gpu' not in args else args['num_gpu']
-                self.GPU_CHIPS = None if 'gpu' not in args else args['gpu']
-                if self.GPU_CHIPS is None:
-                        self.GPU_CHIPS = [0]
-                        if 'CUDA_VISIBLE_DEVICES' in os.environ :
-                                os.environ.pop('CUDA_VISIBLE_DEVICES')
-                        self.NUM_GPUS = 0
-                else:
-                        self.NUM_GPUS = len(self.GPU_CHIPS)
+                # self.GPU_CHIPS = None if 'gpu' not in args else args['gpu']
+                # if self.GPU_CHIPS is None:
+                #         self.GPU_CHIPS = [0]
+                #         if 'CUDA_VISIBLE_DEVICES' in os.environ :
+                #                 os.environ.pop('CUDA_VISIBLE_DEVICES')
+                #         self.NUM_GPUS = 0
+                # else:
+                #         self.NUM_GPUS = len(self.GPU_CHIPS)
                         # os.environ['CUDA_VISIBLE_DEVICES'] = str(self.GPU_CHIPS[0])
-                
+                self.NUM_GPUS = 0 if 'gpu' not in args else args['gpu']
+                self.GPU_CHIPS = None if self.NUM_GPUS == 0 else [args['gpu']]
+                if self.GPU_CHIPS :
+                        os.environ['CUDA_VISIBLE_DEVICES'] = str(self.GPU_CHIPS[0])
                 self.PARTITION = args['partition'] if 'partition' in args else None
                 # if self.NUM_GPUS > 1 :
                 #     os.environ['CUDA_VISIBLE_DEVICES'] = "4"