Hi,
Given the soft implementation of Gumbel Softmax:
from torch.autograd import Variable
import torch.nn.functional as F

def sample_gumbel(input):
    noise = torch.rand(input.size())
    eps = 1e-20
    noise.add_(eps).log_().neg_()
    noise.add_(eps).log_().neg_()
    return Variable(noise)

def gumbel_softmax_sample(input):
    temperature = 1
    noise = sample_gumbel(input)
    x = (input + noise) / temperature
    x = F.log_softmax(x)
    return x.view_as(input)

I’d like to implement the hard version, as here 332.  Any idea how to do that? Maybe this 113 example in torch helps?