sh0416:

The maximum difference between output and pred is just 6.3330e-08 so that the implementation is not wrong but some system issue, I think.


This small difference is most likely a result of the limited floating point precision, which can be seen e.g. by changing the order of operations:
x = torch.randn(10, 10, 10)
sum1 = x.sum()
sum2 = x.sum(0).sum(0).sum(0)

print(sum1 - sum2)
> tensor(-3.8147e-06)

Avoiding these differences is especially hard (or impossible) using different hardware.
What is your current setup, i.e. which GPUs are you using?