Gradients of variables/parameters before softmax becomes zero,
for eg.
import torch
a1 = torch.rand([4, 4], requires_grad=True)
b1 = a1.squeeze(0)**2
b1=F.softmax(b1,dim=1)
b1.sum().backward()
print(a1.grad)
Gives
tensor([[0., 0., 0., 0.],
[0., 0., 0., 0.],
[0., 0., 0., 0.],
[0., 0., 0., 0.]])
How to deal with this in case i need to use softmax in between layers in my model?
Thanks