code: def one_hot(tensor: torch.Tensor, num_classes: int) -> Tensor:
      """Workaround for https://github.com/pytorch/pytorch/issues/55579"""
      assert num_classes > 0, "num_classes must be a positive integer"
      ret = torch.zeros(
          tensor.shape + (num_classes,),
          device=tensor.device,
          dtype=tensor.dtype,
      )
      ret.scatter_(-1, tensor.unsqueeze(-1), 1)
      return ret
  
file_path: fairscale/fairscale/nn/moe/top2gate.py
imports: from typing import Callable, Dict, Tuple
  import torch
  from torch import Tensor
  import torch.nn.functional
label_desc: Use image node_id 2 for calling a global function with example usage: one_hot(tensor, num_classes) and returns: ret