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