Use scatter
def create_one_hot(x: Tensor, n_class: int):
  """
  :param x: [B, D1, D2, D3] 0 <= x.min() && x.max() < C
  :return: [B, C, D1, D2, D3]
  """
  B, D1, D2, D3 = x.shape
  out = torch.zeros((B, n_class, D1, D2, D3), dtype=torch.float, device=x.device)
  x = x[:, None, :, :, :]
  out.scatter_(dim=1, index=x, src=torch.ones_like(x, dtype=out.dtype))
  return out