Hello Blade!



 blade:

Is there a work around for this?


“Work around” isn’t really the right term, as it implies that you are
trying to do something that makes sense, and you need to work
around a bug or limitation to reach the same result by a different
path.
Conceptually, what are the index values in W2 supposed to mean?
How do you want your optimizer to respond when one of the values
in W2 suddenly makes a discrete jump from 2 to 1?
To illustrate my point, here is a “work around”:
Replace W2 = torch.argmax (W1, dim = 3).float()
with W2 = 0.0 * torch.sum (W1, dim = 3).float().
Backpropagation will now work (but all of your gradients will be zero).
softmax() is a smooth (differentiable) approximation to the one-hot
encoding of argmax().  But this comment will only be helpful if you
understand the conceptual role you want the piece-wise constant (not
usefully differentiable) W2 values to play in your network training.
Good luck.
K. Frank