code: def _gumbel_softmax_sample(logits, tau=1, eps=1e-10):
      """
      Draw a sample from the Gumbel-Softmax distribution
  
      based on
      https://github.com/ericjang/gumbel-softmax/blob/3c8584924603869e90ca74ac20a6a03d99a91ef9/Categorical%20VAE.ipynb
      (MIT license)
      """
      if logits.dim() != 3:
          raise AssertionError
      gumbel_noise = _sample_gumbel(logits.size(), eps=eps)
      y = logits + gumbel_noise
      return F.softmax(y / tau, 1)
  
file_path: pywick/pywick/models/segmentation/testnets/gscnn/my_functionals/DualTaskLoss.py
imports: import torch
  import torch.nn
  import torch.nn.functional
  import numpy
  from .custom_functional import compute_grad_mag
label_desc: Use image node_id 3 for calling a global function with example usage: _gumbel_softmax_sample(logits, tau, eps) and returns: F