aou/data-maker
aou
/
data-maker
1
1
Fork 0
Code Issues 4 Pull Requests Packages Projects 1 Releases Wiki Activity

bug fix: gpu visibility

This commit is contained in:
Steve Nyemba 2022-04-12 14:32:39 -05:00
parent d6fd7bceba
commit 838c7978de
1 changed files with 12 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

21
data/gan.py
View File

@ -61,16 +61,19 @@ class GNet :
self.logs = {} self.logs = {}
# self.NUM_GPUS = 1 if 'num_gpu' not in args else args['num_gpu'] # 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'] # self.GPU_CHIPS = None if 'gpu' not in args else args['gpu']
if self.GPU_CHIPS is None: # if self.GPU_CHIPS is None:
self.GPU_CHIPS = [0] # self.GPU_CHIPS = [0]
if 'CUDA_VISIBLE_DEVICES' in os.environ : # if 'CUDA_VISIBLE_DEVICES' in os.environ :
os.environ.pop('CUDA_VISIBLE_DEVICES') # os.environ.pop('CUDA_VISIBLE_DEVICES')
self.NUM_GPUS = 0 # self.NUM_GPUS = 0
else: # else:
self.NUM_GPUS = len(self.GPU_CHIPS) # self.NUM_GPUS = len(self.GPU_CHIPS)
# os.environ['CUDA_VISIBLE_DEVICES'] = str(self.GPU_CHIPS[0]) # 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 self.PARTITION = args['partition'] if 'partition' in args else None
# if self.NUM_GPUS > 1 : # if self.NUM_GPUS > 1 :
# os.environ['CUDA_VISIBLE_DEVICES'] = "4" # os.environ['CUDA_VISIBLE_DEVICES'] = "4"