torch.nn.functional.gumbel_softmax
torch.nn.functional.gumbel_softmax(logits, tau=1, hard=False, eps=1e-10, dim=- 1)
Samples from the Gumbel-Softmax distribution (Link 1  Link 2) and
   optionally discretizes.
Parameters:
      * logits (Tensor) -- [..., num_features] unnormalized
        log probabilities
  * **tau** (*float*) -- non-negative scalar temperature

  * **hard** (*bool*) -- if "True", the returned samples will be
    discretized as one-hot vectors, but will be differentiated as
    if it is the soft sample in autograd

  * **dim** (*int*) -- A dimension along which softmax will be
    computed. Default: -1.

Returns:
      Sampled tensor of same shape as logits from the Gumbel-Softmax
      distribution. If "hard=True", the returned samples will be one-
      hot, otherwise they will be probability distributions that sum
      to 1 across dim.
Return type:
      Tensor
Note: