Your code looks ok from a high level read.
Can you pinpoint when this difference appears? Is it before training? Or one goes down and not the other? Or one diverge and not the other?
PS:
You can initialize your loss accumulator with just python number: total_loss = 0. No need to wrap these.