Solved by god_sp33d in post #4 
              
              
                make the following change in your code in both the places 
net_loss += loss.item()

loss.item() would just sum up the loss value, where as if you add loss it would keep on adding the whole graphs.