From 142491cf5bcd15f61f5cc79023b1cc9dcfeed00f Mon Sep 17 00:00:00 2001 From: Steve Nyemba Date: Wed, 4 Mar 2020 14:27:54 -0600 Subject: [PATCH] Bug Fix: GPU work load --- pipeline.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/pipeline.py b/pipeline.py index a39dbc7..fcf9912 100644 --- a/pipeline.py +++ b/pipeline.py @@ -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