MilesCranmer/backtrack_nans.md

Created March 3, 2021 19:24

Select an option

Learn more about clone URLs
Clone this repository at <script src="https://gist.github.com/MilesCranmer/25442a0058d98b9ae4647439537a41f2.js"></script>
Save MilesCranmer/25442a0058d98b9ae4647439537a41f2 to your computer and use it in GitHub Desktop.

Trick to step backwards if a NaN occurs

Raw

Before training:

last = model.state_dict()

Inside training loop, after computing loss:

if torch.isnan(loss).sum().item():
    model.load_state_dict(last)
else:
    last = model.state_dict()

Select an option