The value of lat is different from total_loss - ce

[JunrQ]

@JunrQ
put lat in loss = [..., stop_grad(lat)]
and print it with print(outputs[-1].asnumpy())