Loaded TemporalFusionTransformer model from S3 bucket cannot make predictions

I used a S3 bucket to store artifacts after training a TemporalFusionTransformer model, and the model itself. Later I needed to load the best model from S3 and make predictions. The model is loaded correctly (the returned object is a TemporalFusionTransformer object) but when I try to make a prediction on that model I get the following:

Traceback (most recent call last):
  File "/home/petartushev/Production script/production_script.py", line 185, in <module>
    preds = model.predict(test_dataloader)
  File "/home/petartushev/Production script/production_venv/lib/python3.7/site-packages/pytorch_forecasting/models/base_model.py", line 1059, in predict
    out = self(x, **kwargs)  # raw output is dictionary
  File "/home/petartushev/Production script/production_venv/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1051, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/petartushev/Production script/production_venv/lib/python3.7/site-packages/pytorch_forecasting/models/temporal_fusion_transformer/__init__.py", line 507, in forward
  File "/home/petartushev/Production script/production_venv/lib/python3.7/site-packages/pytorch_forecasting/models/base_model.py", line 546, in to_network_output
    return self._output_class(**results)
TypeError: __new__() got an unexpected keyword argument 'attention'

When trying to fix this myself, I re-created the model in the same script where I am loading the model, with the same hyperparameters and that re-created model instead of the error above when calling predict with the same test dataloader returned a Tensor as an output.
Any help towards resolving this issue is welcomed.