The error here was because they asked to use return_dict = False but in TF during model compilation, we have to set run_eagerly=True in order to actually make sense of the return_dict parameter.
According to documentation:
return_dict (bool, optional) – Whether or not to return a ModelOutput instead of a plain tuple. This argument can be used in eager mode, in graph mode the value will always be set to True.

Actually it was given in warning that it couldn’t set return_dict to False and hence I was able to figure out this issue.
I don’t know its equivalent for PyTorch but I hope this helps.