code: def one_hot_argmax(inputs, temperature=0.1, axis=-1):
      """Returns one-hot of argmax with backward pass set to softmax-temperature."""
      vocab_size = inputs.shape[-1]
      z = one_hot(torch.argmax(inputs, dim=axis), vocab_size)
      soft = F.softmax(inputs / temperature, dim=axis)
      outputs = soft + (z - soft).detach()
      return outputs
  
file_path: DIG/dig/ggraph/method/GraphDF/model/df_utils.py
imports: import sys
  import warnings
  import torch
  import torch.nn.functional
  import numpy
label_desc: Use image node_id 2 for calling a global function with example usage: one_hot_argmax(inputs, temperature, axis) and returns: outputs