The code snippet works fine in 1.1 but causes an error in 1.2:

p = F.softmax(x, dim=1)
m = y != self.ignore_index
t = F.one_hot((y * m.byte()).long(), num_classes=self.num_classes).byte().permute(0,3,1,2)
i = (p * (t * m.unsqueeze(1).byte()).float()).sum((0,2,3))
u = ((p + t.float()) * m.unsqueeze(1).float()).sum((0,2,3)) - i
v = u.nonzero()
return -((i[v] / u[v]).mean()).log()

I get the error message below:

RuntimeError: range.second - range.first == t.size() INTERNAL ASSERT FAILED at /pytorch/torch/csrc/autograd/generated/Functions.cpp:55, please report a bug to PyTorch. inconsistent range for TensorList output

I had to remove the usage of nonzero() like below to make the code work:

p = F.softmax(x, dim=1)
m = y != self.ignore_index
t = F.one_hot((y * m.byte()).long(), num_classes=self.num_classes).byte().permute(0,3,1,2)
i = (p * (t * m.unsqueeze(1).byte()).float()).sum((0,2,3))
u = ((p + t.float()) * m.unsqueeze(1).float()).sum((0,2,3)) - i
return -((i / u).mean()).log()

I have two questions:
(1) Why does the runtime error happen in 1.2 but not in 1.1?
(2) By not using nonzero to prevent division by zero I have numeric instability in theory. However, in really the number that comes out from a softmax operation should not reach zero because the input can’t really reach negative infinity right?