NHiTS throws an error for dimension mismatch when weight is defined in TimeSeriesDataset

[manitadayon]

• PyTorch-Forecasting version:1.0.0
• PyTorch version: 2.0.0
• Python version: 3.10

Expected behavior

I executed code to be able to handle weight paramter defined in TimeSeriesDataSet, to weight the loss function. However, although the code works for TFT but for NHiTS I got dimension mismatch such as
RuntimeError: The size of tensor a (5) must match the size of tensor b (20) at non-singleton dimension 1

training = TimeSeriesDataSet(
    train_data,
    time_idx="new_time",
    target="var4",
    time_varying_known_categoricals = ['var1','var2','var3'],
    static_categoricals = ['Groups'],
    group_ids=["Groups"],
    time_varying_unknown_reals=["var4'],
    max_encoder_length=context_length,
    max_prediction_length=prediction_length,
    allow_missing_timesteps=False,
    add_target_scales=False,
    )

There is no missing values in my data, since other posts suggest this. Can someone explain how to fix this issue?