I have two code blocks with argmax in both.
First,
x = torch.tensor([0.1, 0.2], dtype=torch.float32, requires_grad=True)
y = torch.argmax(x, dim=-1)
y = torch.unsqueeze(y, 0)
y = y.to(torch.float32)
y = nn.Linear(1, 1)(y)
y.backward()

And second,
x = torch.tensor([0.1, 0.2], dtype=torch.float32, requires_grad=True)
y = torch.argmax(x, dim=-1)
y.backward()

Interesting thing is, the first code can run successful, but the second failed. So why?