Sure.
model_global.to(device) is a typo error it would be model_net.to(device). I rewrote them.
I want to use grad to sending it to other virtual machines to update the models there. Moreover, with loss.backward(retain_graph=True) the model would update in terms of all leaf nodes (data and weights).
The following part is just using if I would pass some data to be trained not all the data. However, in the model () function I want to pass all data so the following part is not needed.
if iteration_count >= max_nr_batches >= 0:
    break