def _cast_types(self, input_dict):
"""This function performs automatic cast of all inputs to encoder dtype.
Args:
input_dict (dict): dictionary passed to :meth:`self._encode() <_encode>`
method.
Returns:
dict: same as input_dict, but with all Tensors cast to encoder dtype.
"""
return cast_types(input_dict, self.params['dtype'])
def _cast_types(self, input_dict):
"""This function performs automatic cast of all inputs to the loss dtype.
Args:
input_dict (dict): dictionary passed to
:meth:`self._compute_loss() <_compute_loss>` method.
Returns:
dict: same as input_dict, but with all Tensors cast to the loss dtype.
"""
return cast_types(input_dict, self.params['dtype'])