Change learning_rate and max_steps after loading models

[jenaparadies]

Hi,

it was previously discussed here: #988
I can't find a way to change e.g. max_steps in the models after they've been set up or loaded. As you can see in the example below, the models are built with max_steps=500, and immediately after I want to change it to 100. But the models are all trained with max_steps=500. This is particularly bad after saving and loading a model, since the learning_rate is also baked into the model, so if a model is stuck, you have to start from scratch.

Here is the example from the tutorial:

from neuralforecast.utils import AirPassengersDF

Y_df = AirPassengersDF
Y_df.head()

import logging

from neuralforecast import NeuralForecast
from neuralforecast.models import LSTM, NHITS, RNN, TimesNet

logging.getLogger('pytorch_lightning').setLevel(logging.ERROR)

horizon = 12

# Try different hyperparmeters to improve accuracy.
models = [LSTM(input_size=2 * horizon,
               h=horizon,                    # Forecast horizon
               max_steps=500,                # Number of steps to train
               scaler_type='standard',       # Type of scaler to normalize data
               encoder_hidden_size=64,       # Defines the size of the hidden state of the LSTM
               decoder_hidden_size=64,),     # Defines the number of hidden units of each layer of the MLP decoder
          NHITS(h=horizon,                   # Forecast horizon
                input_size=2 * horizon,      # Length of input sequence
                max_steps=500,               # Number of steps to train
                n_freq_downsample=[2, 1, 1]), # Downsampling factors for each stack output
          TimesNet(h=horizon,                   # Forecast horizon
                input_size=2 * horizon,      # Length of input sequence
                max_steps=500,)               # Number of steps to train
          ]


nf = NeuralForecast(models=models, freq='ME')

#test if we can change model parameters --> it doesn't work, nf.fit will use max_steps=500 from the initialisation
nf.models[0].max_steps = 100
nf.models[1].max_steps = 100
nf.models[2].max_steps = 100

nf.fit(df=Y_df)

Y_hat_df = nf.predict()

print(Y_hat_df.head())

import matplotlib.pyplot as plt

plt.figure(figsize=(10, 5))
plt.plot(Y_df['ds'], Y_df['y'], label='Ground Truth', linestyle='--')
plt.plot(Y_hat_df['ds'], Y_hat_df['TimesNet'], label='TimesNet')#, marker='o')
plt.plot(Y_hat_df['ds'], Y_hat_df['NHITS'], label='NHITS')#, marker='o')
plt.plot(Y_hat_df['ds'], Y_hat_df['LSTM'], label='LSTM')#, marker='o')

plt.title(" Forecast with TimesNet vs NHITS vs LSTM")
plt.xlabel("Date")
plt.ylabel("Value")
#plt.yscale("log")

plt.legend()
plt.grid(True)

plt.show()