Pytorch N - Beats model throwing error: 'str' object has no attribute '__name__'

I'm trying to replicate pytorch's N - Beats model in colab. I copied the same code from https://pytorch-forecasting.readthedocs.io/en/stable/tutorials/ar.html to a colab notebook. There is an error showing up at training cell.

import os 
import warnings
warnings.filterwarnings("ignore")
os.chdir("../../..")

import pandas as pd
import pytorch_lightning as pl
from pytorch_lightning.callbacks import EarlyStopping
import torch

from pytorch_forecasting import Baseline, NBeats, TimeSeriesDataSet
from pytorch_forecasting.data import NaNLabelEncoder
from pytorch_forecasting.data.examples import generate_ar_data
from pytorch_forecasting.metrics import SMAPE

data = generate_ar_data(seasonality=10.0, timesteps=400, n_series=100, seed 
 = 42)
data["static"] = 2
data["date"] = pd.Timestamp("2020-01-01") + pd.to_timedelta(data.time_idx, "D")
data.head()

# create dataset and dataloaders
max_encoder_length = 60
max_prediction_length = 20

training_cutoff = data["time_idx"].max() - max_prediction_length

context_length = max_encoder_length
prediction_length = max_prediction_length

training = TimeSeriesDataSet(
data[lambda x: x.time_idx <= training_cutoff],
time_idx="time_idx",
target="value",
categorical_encoders={"series": NaNLabelEncoder().fit(data.series)},
group_ids=["series"],
# only unknown variable is "value" - and N-Beats can also not take any additional variables
 time_varying_unknown_reals=["value"],
 max_encoder_length=context_length,
 max_prediction_length=prediction_length,
)

validation = TimeSeriesDataSet.from_dataset(training, data,     min_prediction_idx=training_cutoff + 1)
batch_size = 128
train_dataloader = training.to_dataloader(train=True,     batch_size=batch_size, num_workers=0)
val_dataloader = validation.to_dataloader(train=False, batch_size=batch_size, num_workers=0)

Error is:

---------------------------------------------------------------------------

AttributeError                            Traceback (most recent call last)

<ipython-input-67-db4b0ef13391> in <module>()
     25     net,
     26     train_dataloader=train_dataloader,
---> 27     val_dataloaders=val_dataloader,
     28 )

30 frames

/usr/local/lib/python3.7/dist-packages/yaml/representer.py in     represent_object(self, data)
    329         if dictitems is not None:
    330             dictitems = dict(dictitems)
--> 331         if function.__name__ == '__newobj__':
    332             function = args[0]
    333             args = args[1:]

 AttributeError: 'str' object has no attribute '__name__'