[FIX] Multivariate models give error when predicting when n_series > batch_size

[elephaint]

Multivariate models give an error when predicting when n_series > batch_size. I covered this for training, but not for predicting. Subtle bug, easy fix.

The below code fails without the fix, and runs correctly with it:

import pandas as pd
import matplotlib.pyplot as plt

from neuralforecast import NeuralForecast
from neuralforecast.models import TSMixerx, NHITS
from neuralforecast.utils import generate_series

N_SERIES = 1025
FREQ= 'D'
df = generate_series(n_series=N_SERIES, seed=0, freq=FREQ, equal_ends=True)
max_ds = df.ds.max() - pd.Timedelta(14, FREQ)
Y_TRAIN_DF = df[df.ds < max_ds]
Y_TEST_DF = df[df.ds >= max_ds]

models = [TSMixerx(h=12,
                input_size=24,
                n_series=N_SERIES,
                max_steps=10,
                batch_size=10,
                revin=True,
                valid_batch_size=10,
                ),
            NHITS(h=12,
                input_size=24,
                max_steps=10,
                batch_size=10,
                valid_batch_size=100,
                )
            ]
        

fcst = NeuralForecast(models=models, freq=FREQ)
fcst.fit(df=Y_TRAIN_DF)
forecasts = fcst.predict(futr_df=Y_TEST_DF)

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[marcopeix]

LGTM