diff --git a/pipeline.py b/pipeline.py
index 0583116..f32a45e 100644
--- a/pipeline.py
+++ b/pipeline.py
@@ -142,7 +142,7 @@ class Components :
 		#
 		# We need to make sure that continuous columns are removed 
 		if x_cols :
-			_args['data'] = _args['data'][list(set(df.columns) - set(x_cols))]
+			_args['data'] = _args['data'][list(set(_args['data'].columns) - set(x_cols))]
 		if 'gpu' in args :
 			_args['gpu'] = self.set_gpu(gpu=args['gpu'])
 		if df.shape[0] and df.shape[0] :
@@ -237,7 +237,7 @@ class Components :
 		# @TODO: Abstract this !!
 		#
 		if x_cols :
-			args['data'] = args['data'][list(set(df.columns) - set(x_cols))]
+			args['data'] = args['data'][list(set(args['data'].columns) - set(x_cols))]
 		
 		args['candidates']	= 1 if 'candidates' not in args else int(args['candidates'])
 		if 'gpu' in args :