about a problem in gpu_version

[Astrosolar]

thanks your codes.
I finished the normal version, but when I used the dist_version, it occurs a issue.
The errors shows "AssertionError :assert not np.isnan(loss_value)" in distributed_train.py. in line 132.
how's it going? can you help me ?