diff --git a/forecastout/core.py b/forecastout/core.py index 594a8d2..e92a7c1 100644 --- a/forecastout/core.py +++ b/forecastout/core.py @@ -157,7 +157,9 @@ def __get_daily_disaggregation_shares(self): ) self.log.info(">> 3.3. Train disaggregation models") disaggregation_model_trainer = DisaggregationModelTrainer( - df_train=df_train_test.loc[df_train_test['test'] == 0].copy(), + df_train=( + df_train_test.loc[df_train_test['test'] == 0].dropna().copy() + ), config=self.config ) disaggregation_model_trained = disaggregation_model_trainer.train(