import torch


def argSoftmax1d(tensor, temperature):
    """
  Find the index of the Softmax value in a 1d tensor. Prob is proportional to q/temperature

  Args:
    - tensor: PyTorch tensor of size (n x 1 x d)

  Returns: nx1 PyTorch tensor containing indexes of max values
  """
    n = tensor.size(0)
    probs = (tensor / temperature).view(n, -1).softmax(dim=-1)
    m = torch.multinomial(probs, 1)
    return m.long()