Only extract one word from gumbel softmax

[allanj]

In the code https://github.com/yala/text_nn/blob/master/rationale_net/utils/learn.py#L71-L85

def get_hard_mask(z, return_ind=False):
    '''
        -z: torch Tensor where each element probablity of element
        being selected
        -args: experiment level config
        returns: A torch variable that is binary mask of z >= .5
    '''
    max_z, ind = torch.max(z, dim=-1)
    if return_ind:
        del z
        return ind
    masked = torch.ge(z, max_z.unsqueeze(-1)).float()
    del z
    return masked

because we take the max, usually, only one position will have the max value.
In this case, if we have 100 words in the sentence, we only select one word as the rationale?
I thought we should select independently and choose those words with >0.5 probability.

Maybe we should change

masked = torch.ge(z, max_z.unsqueeze(-1)).float()

to

masked = torch.ge(z, 0.5).float()

instead?