Gracefully shutdown when loss is NaN

[Arcomano1234]

We raise an exception when the loss is NaN, where this check is done locally on each rank (see _validate_loss in fme/core/optimization.py). However, there is a chance that there is only an NaN loss on a subset of the ranks. When this error is not raised on all the ranks the other ranks call torch.backward, which is a collective all reduce and instead the other GPUs hang until a NCCL timeout error occurs.

Changes:

• add dist shutdown to _validate_loss in fme/core/optimization.py if the training is distributed

• Tests added

Resolves #849