bug fix with GPU allocation

This commit is contained in:
Steve Nyemba 2020-03-09 13:10:26 -05:00
parent 8455cd7554
commit bbd03c4a63
1 changed files with 3 additions and 3 deletions

View File

@ -84,7 +84,7 @@ class Components :
# We ask the process to assume 1 gpu given the system number of GPU and that these tasks can run in parallel # We ask the process to assume 1 gpu given the system number of GPU and that these tasks can run in parallel
# #
if int(args['num_gpu']) > 1 : if int(args['num_gpu']) > 1 :
_args['gpu'] = int(args['gpu']) if int(args['gpu']) < 8 else np.random.choice(np.arange(8)).astype(int)[0] _args['gpu'] = int(args['gpu']) if int(args['gpu']) < 8 else np.random.choice(np.arange(8)).astype(int)
else: else:
_args['gpu'] = 0 _args['gpu'] = 0
_args['num_gpu'] = 1 _args['num_gpu'] = 1
@ -124,7 +124,7 @@ class Components :
_args['batch_size'] = int(args['batch_size']) _args['batch_size'] = int(args['batch_size'])
if int(args['num_gpu']) > 1 : if int(args['num_gpu']) > 1 :
_args['gpu'] = int(args['gpu']) if int(args['gpu']) < 8 else np.random.choice(np.arange(8)).astype(int)[0] _args['gpu'] = int(args['gpu']) if int(args['gpu']) < 8 else np.random.choice(np.arange(8)).astype(int)
else: else:
_args['gpu'] = 0 _args['gpu'] = 0
_args['num_gpu'] = 1 _args['num_gpu'] = 1
@ -215,7 +215,7 @@ class Components :
df = pd.DataFrame(info['data']) df = pd.DataFrame(info['data'])
args = info['args'] args = info['args']
if args['num_gpu'] > 1 : if args['num_gpu'] > 1 :
args['gpu'] = int(info['input']['partition']) if info['input']['partition'] < 8 else np.random.choice(np.arange(8),1).astype(int)[0] args['gpu'] = int(info['input']['partition']) if info['input']['partition'] < 8 else np.random.choice(np.arange(8)).astype(int)
else: else:
args['gpu'] = 0 args['gpu'] = 0