Fixing unit test

onnx · May 21, 2022 · bcabd14 · bcabd14
1 parent 446b7a9
commit bcabd14
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/sparkml/test_pipeline.py b/tests/sparkml/test_pipeline.py
@@ -125,7 +125,7 @@ def test_model_pipeline_2_stage(self):
         stages = []
         for col in cols:
             stages.append(StringIndexer(inputCol=col, outputCol=col+'_index', handleInvalid='skip'))
-            stages.append(OneHotEncoder(inputCols=[col+'_index'], outputCols=[col+'_vec']))
+            stages.append(OneHotEncoder(inputCols=[col+'_index'], outputCols=[col+'_vec'], dropLast=False))
 
         pipeline = Pipeline(stages=stages)