This commit is contained in:
Steve Nyemba 2022-04-12 14:50:19 -05:00
parent 4aaefedce0
commit bbbeb5172a
2 changed files with 11 additions and 13 deletions

View File

@ -61,19 +61,16 @@ 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"

View File

@ -319,7 +319,8 @@ class Generator (Learner):
_args['map'] = self._map _args['map'] = self._map
_args['values'] = np.array(values) _args['values'] = np.array(values)
_args['row_count'] = self._df.shape[0] _args['row_count'] = self._df.shape[0]
if self.gpu :
_args['gpu'] = self.gpu
gHandler = gan.Predict(**_args) gHandler = gan.Predict(**_args)
gHandler.load_meta(columns=None) gHandler.load_meta(columns=None)
_iomatrix = gHandler.apply() _iomatrix = gHandler.apply()