From bbd03c4a63aeb109dd878d4e50cd9b8568bf8b45 Mon Sep 17 00:00:00 2001 From: Steve Nyemba Date: Mon, 9 Mar 2020 13:10:26 -0500 Subject: [PATCH] bug fix with GPU allocation --- pipeline.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pipeline.py b/pipeline.py index 89ba16f..7a2cf3a 100644 --- a/pipeline.py +++ b/pipeline.py @@ -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 # 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: _args['gpu'] = 0 _args['num_gpu'] = 1 @@ -124,7 +124,7 @@ class Components : _args['batch_size'] = int(args['batch_size']) 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: _args['gpu'] = 0 _args['num_gpu'] = 1 @@ -215,7 +215,7 @@ class Components : df = pd.DataFrame(info['data']) args = info['args'] 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: args['gpu'] = 0