Bug Fix: GPU work load

This commit is contained in:
Steve Nyemba 2020-03-04 14:27:54 -06:00
parent 2b7b1757f9
commit 142491cf5b
1 changed files with 4 additions and 3 deletions

View File

@ -209,9 +209,10 @@ class Components :
logger.write({'module':'process','action':'read-partition','input':info['info']})
df = pd.DataFrame(info['data'])
args = info['args']
args['gpu'] = int(info['info']['partition'])
if int(args['num_gpu']) > 1 and args['gpu'] > 0:
args['gpu'] = args['gpu'] + args['num_gpu'] if args['gpu'] + args['num_gpu'] < 8 else 0 #-- 8 max gpus
MAX_GPUS = 8
args['gpu'] = int(info['info']['partition']) if info['info']['partition'] < MAX_GPUS else np.random.choice(np.arange(MAX_GPUS),1).astype(int).tolist()[0]
# if int(args['num_gpu']) > 1 and args['gpu'] > 0:
# args['gpu'] = args['gpu'] + args['num_gpu'] if args['gpu'] + args['num_gpu'] < 8 else args['gpu'] #-- 8 max gpus
args['reader'] = lambda: df
#
# @TODO: Fix