[AutoPatchTST] TypeError: Trainer.__init__() got an unexpected keyword argument 'n_head'

[ams015]

What happened + What you expected to happen

When doing hyperparameter selection with AutoPatchTST,if I set the config=None, I am hoping for it to work successfully still, like it does with other models.

File "/usr/local/lib/python3.10/dist-packages/ray/air/execution/_internal/event_manager.py", line 110, in resolve_future
result = ray.get(future)
File "/usr/local/lib/python3.10/dist-packages/ray/_private/auto_init_hook.py", line 24, in auto_init_wrapper
return fn(*args, **kwargs)
File "/usr/local/lib/python3.10/dist-packages/ray/_private/client_mode_hook.py", line 103, in wrapper
return func(*args, **kwargs)
File "/usr/local/lib/python3.10/dist-packages/ray/_private/worker.py", line 2563, in get
raise value.as_instanceof_cause()
ray.exceptions.RayTaskError(TypeError): ray::ImplicitFunc.train() (pid=2265, ip=172.28.0.12, actor_id=6b35db0c8eb04cf8f2da426701000000, repr=_train_tune)
File "/usr/local/lib/python3.10/dist-packages/ray/tune/trainable/trainable.py", line 342, in train
raise skipped from exception_cause(skipped)
File "/usr/local/lib/python3.10/dist-packages/ray/air/_internal/util.py", line 91, in run
self._ret = self._target(*self._args, **self._kwargs)
File "/usr/local/lib/python3.10/dist-packages/ray/tune/trainable/function_trainable.py", line 115, in
training_func=lambda: self._trainable_func(self.config),
File "/usr/local/lib/python3.10/dist-packages/ray/tune/trainable/function_trainable.py", line 332, in _trainable_func
output = fn()
File "/usr/local/lib/python3.10/dist-packages/ray/tune/trainable/util.py", line 321, in inner
return trainable(config, **fn_kwargs)
File "/usr/local/lib/python3.10/dist-packages/neuralforecast/common/_base_auto.py", line 207, in _train_tune
_ = self._fit_model(
File "/usr/local/lib/python3.10/dist-packages/neuralforecast/common/_base_auto.py", line 336, in _fit_model
model.fit(dataset, val_size=val_size, test_size=test_size)
File "/usr/local/lib/python3.10/dist-packages/neuralforecast/common/_base_windows.py", line 733, in fit
trainer = pl.Trainer(**self.trainer_kwargs)
File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/utilities/argparse.py", line 70, in insert_env_defaults
return fn(self, **kwargs)
TypeError: Trainer.init() got an unexpected keyword argument 'n_head'

Versions / Dependencies

Python 3.10.12, neuralforecast 1.6.4, on google colab

Reproduction script

from neuralforecast.tsdataset import TimeSeriesDataset
from neuralforecast.utils import AirPassengers, AirPassengersPanel, AirPassengersStatic
import pandas as pd
from neuralforecast import NeuralForecast
from ray import tune
from neuralforecast.auto import AutoPatchTST
from neuralforecast.losses.pytorch import MAE

#AirPassengersPanel['y'] = AirPassengersPanel['y'] + 10
Y_train_df = AirPassengersPanel[AirPassengersPanel.ds<AirPassengersPanel['ds'].values[-12]] # 132 train
Y_test_df = AirPassengersPanel[AirPassengersPanel.ds>=AirPassengersPanel['ds'].values[-12]].reset_index(drop=True) # 12 test

models2= [AutoPatchTST(h=15,
loss=MAE(),
config=None,
num_samples=2,gpus=1 )]
nf= NeuralForecast(
models=models2,
freq='b')
nf.fit(df=Y_train_df)

Issue Severity

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

[cchallu]

Hi @ams015! Thanks for noting the error. There was an error with the name of the n_head parameter, it should be n_heads. I just created a PR to fix the issue, it will be merged to the main branch soon. While we create the new release, you can install the main branch in colab with pip install git+https://github.com/Nixtla/neuralforecast.git@main