[BUG] ETSformer training on GPU

[Sharaddition]

ETSformerForecaster trainer does not work on GPU notebooks.

Error:

indices should be either on cpu or on the same device as the indexed tensor (cpu)

To Reproduce

# Taken from Dashboard code
params = {'max_forecast_steps': forecast_steps,'n_past': n_past, 'use_gpu': True}
model_class = ModelFactory.get_model_class('ETSformerForecaster')
model = model_class(model_class.config_class(**params))

# EXOG DATA
if model.supports_exog and len(exog_columns) > 0:
    print('Exog Support: True')
    exog_ts = TimeSeries.from_pd(pd.concat((train_df.loc[:, exog_columns], test_df.loc[:, exog_columns])))
    train_df = train_df.loc[:, [target_column] + feature_columns]
    test_df = test_df.loc[:, [target_column] + feature_columns]
else:
    print('Exog Support: False')
    exog_ts = None

train_ts = TimeSeries.from_pd(train_df)
predictions = model.train(train_ts, exog_data=exog_ts)

Device:

• Colab/Kaggle Notebook
• Merlion Version 2.0.2

Additional context
Model trains perfectly on CPU, but as expected takes a long time to train on large data.

Screenshot