def METHOD_NAME(self, spatial_input, non_spatial_input, action_mask):
    values, actions = self(spatial_input, non_spatial_input)
    # Masking step: Inspired by: http://juditacs.github.io/2018/12/27/masked-attention.html
    if action_mask is not None:
        actions[~action_mask] = float('-inf')
    action_probs = F.softmax(actions, dim=1)
    return values, action_probs