def METHOD_NAME(
    self,
    inputs: Union[
        Dict[str, torch.Tensor], Dict[str, np.ndarray], Tuple[Dict[str, torch.Tensor], Dict[str, torch.Tensor]]
    ],
    mask=None,
) -> Dict[str, torch.Tensor]:
    """Forward pass of the model.
    Args:
        inputs: Inputs to the model. Can be a dictionary of input names to
            input tensors or a tuple of (inputs, targets) where inputs is
            a dictionary of input names to input tensors and targets is a
            dictionary of target names to target tensors.
        mask: A mask for the inputs.
    Returns:
        A dictionary of output {feature name}::{tensor_name} -> output tensor.
    """
    if isinstance(inputs, tuple):
        inputs, targets = inputs
        # Convert targets to tensors.
        for target_feature_name, target_value in targets.items():
            if not isinstance(target_value, torch.Tensor):
                targets[target_feature_name] = torch.from_numpy(target_value)
            else:
                targets[target_feature_name] = target_value
    else:
        targets = None
    assert list(inputs.keys()) == self.input_features.keys()
    encoder_outputs = self.encode(inputs)
    combiner_outputs = self.combine(encoder_outputs)
    return self.decode(combiner_outputs, targets, mask)