diff options
Diffstat (limited to 'ml/model_trainer.py')
| -rw-r--r-- | ml/model_trainer.py | 5 |
1 files changed, 3 insertions, 2 deletions
diff --git a/ml/model_trainer.py b/ml/model_trainer.py index acd7218..79ce706 100644 --- a/ml/model_trainer.py +++ b/ml/model_trainer.py @@ -12,6 +12,7 @@ import random import sys from types import SimpleNamespace from typing import Any, List, NamedTuple, Optional, Set, Tuple +import warnings import numpy as np from sklearn.model_selection import train_test_split # type:ignore @@ -247,13 +248,13 @@ class TrainingBlueprint(ABC): if self.spec.delete_bad_inputs: msg = f"WARNING: {filename}: missing features or label; expected {self.spec.feature_count} but saw {len(x)}. DELETING." - print(msg, file=sys.stderr) logger.warning(msg) + warnings.warn(msg) os.remove(filename) else: msg = f"WARNING: {filename}: missing features or label; expected {self.spec.feature_count} but saw {len(x)}. Skipping." - print(msg, file=sys.stderr) logger.warning(msg) + warnings.warn(msg) return (X, y) def make_progress_graph(self) -> None: |