Hi Trick!



 trick:

y = torch.argmax(x, dim=-1)



As a side note, at this point y is  long so it can’t have
requires_grad = True.  So you wouldn’t be able to run
.backward() on it.


y = y.to(torch.float32)
y = nn.Linear(1, 1)(y)
y.backward()



Here, you make y a float (but it still doesn’t have
requires_grad = True).  Then when you run it through Linear,
it gets multiplied my Linear's weight which does have
requires_grad = True (and gets bias added to it which also has
requires_grad = True).  So the result has requires_grad = True,
and you can run .backward() on it.


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



Here, the fact that x has requires_grad = True is irrelevant.  The
result of argmax() is  long (which can’t have requires_grad = True),
so you can’t run .backward() on it.
Best.
K. Frank