NHITS does not give consistent results on GPU

[Mickailkhadhar]

What happened + What you expected to happen

Hello, I'm tuning NHITS model using Optuna when I discovered I had inconsistent results with the same experiments.
I replicated the bug using utilsforecast to generate series.
What I do in my code is generating series once. Then I create the exact same model twice.
Predictions are always different using a GPU on Databricks.
It seems that the bug never occurs when using CPU in Databricks too. I also tested it with neuralforecast 1.7.3 and 1.7.5 version.

Predictions seem inconsistent regardless of the hyperparameters values and for equal random_seed.
PS : I am also using NBEATSx and TFT and have no problem with both!

Thanks in advance, Appreciate your help!

Versions / Dependencies

Within Databricks environment.
Using neuralforecast 1.7.3 and 1.75
Bug reproducible on GPU but not on CPU

Reproduction script

from neuralforecast import NeuralForecast
from neuralforecast.models import NHITS
from utilsforecast.data import generate_series


n_series = 1500
n_static_features = 10

df = generate_series(
    n_series=n_series,
    freq="W",
    min_length=208,
    max_length=208,
    equal_ends=True,
    n_static_features=n_static_features,
)

# Split the data into train and test
h = 52
train_df = df.groupby('unique_id').apply(lambda x: x.iloc[:-h]).reset_index(drop=True)
test_df = df.groupby('unique_id').apply(lambda x: x.iloc[-h:]).reset_index(drop=True)

# Create static features dataframe
static_df = df.groupby("unique_id", as_index=False).first().drop(columns=["y", "ds"])

## Modeling 

model = NHITS(h=h,
            input_size=104,
            stat_exog_list=[f"static_{i}" for i in range(n_static_features)],
            activation="PReLU", 
            n_blocks=[3, 3, 3],
            mlp_units=4*[[256, 256]],
            max_steps=100,
            random_seed=42,
        )

model2 = NHITS(h=h,
            input_size=104,
            stat_exog_list=[f"static_{i}" for i in range(n_static_features)],
            activation="PReLU", 
            n_blocks=[3, 3, 3],
            mlp_units=4*[[256, 256]],
            max_steps=100,
            random_seed=42,
        )

# Initialize NeuralForecast object
nf = NeuralForecast(models=[model], freq='W-SUN')
nf2 = NeuralForecast(models=[model2], freq='W-SUN')

# Fit the model on the training data
nf.fit(df=train_df, static_df=static_df)
nf2.fit(df=train_df, static_df=static_df)

# Predict using the test set
predictions = nf.predict(df=train_df, static_df=static_df, futr_df=test_df)
predictions2 = nf2.predict(df=train_df, static_df=static_df, futr_df=test_df)

# Comparing both predictions, bug should output 78000

(predictions == predictions2)['NHITS'].value_counts()[False]

Issue Severity

High: It blocks me from completing my task.