we could carry following experiment to see what is going on while computing gradient, let us assume we have
l = nn.Linear(2, 2, bias=False)
list(l.parameters())

[Parameter containing:
tensor([[-0.0473,  0.0679],
[-0.6429, -0.5586]], requires_grad=True)]
input = torch.tensor([[1.3032, 0.5701], [0.2671, 1.3900]])

o1 = l(input)
o1

tensor([[-0.0230, -1.1563],
[ 0.0817, -0.9482]], grad_fn=<MmBackward>)
loss_fn = nn.CrossEntropyLoss(reduction='sum')
loss = loss_fn(o1, torch.tensor([0, 1]))
# assume that target is [0, 1]

softmax (nn.Softmax(dim=-1)) is,

Screenshot (248)1808×334 10.3 KB

output after passing input through linear and then softmax would be like,
(w00, w01, w10, w11 are weights)

Screenshot (251)1780×266 22.7 KB

after passing this through cross entropy loss with reduction=‘sum’, (we do not apply softmax externally, cross entropy has inbuilt softmax),

Screenshot (253)1738×229 16.2 KB

(since our target was torch.tensor([0, 1]), we considered 0th index in first row, and 1st index in second row)
after taking jacobian and substituting weights with their values,

Screenshot (276)1795×154 8.58 KB


Screenshot (256)1659×640 46.1 KB

if we do
loss.backward()
for param in l.parameters():
  print(param.grad)

then we get,
tensor([[-0.1206,  0.8855],
[ 0.1206, -0.8855]])
which matches with our calculated gradient.
if our output after applying linear contained float('-inf'), that is it was something like,
o1 = torch.tensor([[float('-inf'), 2.], [1., 2.]], requires_grad=True)

then after applying cross entropy,

Screenshot (270)1778×284 12 KB

after computing jacobian,

Screenshot (272)1517×534 14.6 KB

substitute with values,

Screenshot (274)1698×327 16.5 KB

which matches with,
o1 = torch.tensor([[float('-inf'), 2.], [1., 2.]], requires_grad=True)
loss_fn = nn.CrossEntropyLoss(reduction='sum')
loss = loss_fn(o1, torch.tensor([0, 1]))
loss.backward()
o1.grad

tensor([[-1.0000,  1.0000],
[ 0.2689, -0.2689]])
all the float(’-inf’) terms vanish