gitea_admin/sientia-dataops-model-manager
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: clear input example in Training class before model storage

Browse Source 
- 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
1 changed files with 1 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
model_manager/activities/training.py
View File

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