[TSMixerx] TypeError: Trainer.__init__() got an unexpected keyword argument 'optimizer'

[RoelofKuipers]

What happened + What you expected to happen

I tried to run to use the TSMixer.py class, but got an error :
TypeError: Trainer.__init__() got an unexpected keyword argument 'optimizer'
I think it is because because of the pytorch-lightning does not have an optimizer argument in the Trainer constructor anymore. This applies to these two lines in tsmixerx.py

Versions / Dependencies

pytorch-lightning=2.0.0
torch==2.1.0
neuralforecast==1.6.4

Reproduction script

import numpy as np
import pandas as pd
import pytorch_lightning as pl
import matplotlib.pyplot as plt

from neuralforecast import NeuralForecast
from neuralforecast.utils import AirPassengersPanel, AirPassengersStatic
from neuralforecast.losses.pytorch import MAE
from neuralforecast.models import TSMixerx 

Y_train_df = AirPassengersPanel

horizon = 12
model = TSMixerx(h=horizon,
                input_size=24,
                n_series=2,
                hist_exog_list=['trend','y_[lag12]'],
                #futr_exog_list=['trend'],
                n_block=4,
                ff_dim=4,
                revin=True,
                scaler_type='standard',
                max_steps=1000,
                early_stop_patience_steps=3,
                val_check_steps=50,
                learning_rate=1e-3,
                loss=MAE(),
                valid_loss=MAE(),
                batch_size=32
                )

# We make validation and test splits
n_time = len(Y_train_df.ds.unique())
val_size = int(.2 * n_time)
test_size = int(.2 * n_time)

fcst = NeuralForecast(models=[model], freq='M')

Y_hat_df = fcst.cross_validation(df=Y_train_df, val_size=val_size,
                               test_size=test_size, n_windows=None)
     
Y_hat_df = Y_hat_df.reset_index()     

Issue Severity

Medium: It is a significant difficulty but I can work around it.

[jmoralez]

Hey @RoelofKuipers, thanks for using neuralforecast. That argument refers to the feature added in #901, so it's not present in 1.6.4. If you want to try the new models I suggest installing from github instead, so that you get a compatible version.

[JQGoh]

@RoelofKuipers I have just made a quick test using the dev env for TSMixerx and it works fine with the optimizer argument. Hope this helps

[RoelofKuipers]

Thank you :)