I am also interested in implementing ST gumbel softmax with stop gradient.
So, I wrote the script as below.
def gumbel_softmax_sample(self, input):
        temperature = 1
        noise = self.sample_gumbel(input)
        x = (input + noise) / temperature
        x = F.softmax(x)

        if hard == True:
            max_val, _ = torch.max(x, x.dim()-1)
            x_hard = x == max_val.expand_as(x)
            tmp = (x_hard.float() - x)
            tmp2 = tmp.clone()
            tmp2.detach_()
            x = tmp2 + x

        return x.view_as(input)

But, I tried to modify it such that only one ‘1’ exists at each row in x_hard as below.
def gumbel_softmax_sample(self, input):
    temperature = 1
    noise = self.sample_gumbel(input)
    x = (input + noise) / temperature
    x = F.softmax(x)

    if hard == True:
        _, max_inx = torch.max(x, x.dim()-1)
        x_hard = torch.cuda.FloatTensor(x.size()).zero_().scatter_(x.dim()-1, max_inx.data, 1.0)
        x2 = x.clone()
        tmp = Variable(x_hard-x2.data)
        tmp.detach_()

        x = tmp + x
        
    return x.view_as(input)

These two scripts work without error, but I am not sure they are equivalent to the original ST gumbel softmax.
I would greatly appreciate if someone reviews them.