diff --git a/app/preprocessing/transform_dataset.py b/app/preprocessing/transform_dataset.py index 5711f94a9cc03e5a27b2767eede983db5acde0ee..914be5ad0cb0839ebcf8a132fbf960bf3ee4dd11 100644 --- a/app/preprocessing/transform_dataset.py +++ b/app/preprocessing/transform_dataset.py @@ -194,7 +194,7 @@ def _remove_unimportant_predictors(train_files: List[Path], all_files: List[Path df = pd.read_parquet(file) columns_to_drop = [col for col in df.columns if col not in columns_to_keep] - df.drop(columns=columns_to_drop) + df = df.drop(columns=columns_to_drop) df.to_parquet(out_dir / file.name)