fix: clear input example in Training class before model storage

- Set wrapper._input_example to None to ensure no residual data is stored with the model, enhancing data integrity during model storage.
This commit is contained in:
vitor-aignosi
2026-04-16 15:08:26 -03:00
parent e2b2f702f8
commit b245801e09

View File

@@ -338,6 +338,7 @@ class Training(SientiaMonitoring):
raise ValueError('Report path, train data path, or test data path is not set')
self.info(f'Storing model for {train_params.model_type}', metadata)
wrapper._input_example = None
wrapper.store_model(name=train_params.model_name)
self.info(f'Logging artifacts for {train_params.model_type}', metadata)
mlflow.log_artifact(train_result.report_path)