drop sampling key while crossvalidating

torronen · Aug 8, 2023 · 076fd26 · 076fd26
1 parent 077a6b8
commit 076fd26
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/src/Microsoft.ML.Data/TrainCatalog.cs b/src/Microsoft.ML.Data/TrainCatalog.cs
@@ -102,7 +102,10 @@ internal CrossValidationResult(ITransformer model, T metrics, IDataView scores,
             // spawn off a separate host per fold in that case.
             foreach (var split in DataOperationsCatalog.CrossValidationSplit(Environment, data, splitColumn, numFolds))
             {
-                var model = estimator.Fit(split.TrainSet);
+                var trainSetWithoutSamplingKey = ColumnSelectingTransformer.CreateDrop(Environment, split.TrainSet, samplingKeyColumn);
+                var testSetWithoutSamplingKey = ColumnSelectingTransformer.CreateDrop(Environment, split.TestSet, samplingKeyColumn);
+
+                var model = estimator.Fit(trainSetWithoutSamplingKey);
                 IDataView scoredTest;
 
                 if (IsCastableToTransformerChainOfITransformer(model))