I tried again with the saving loading of the optimizer as well, and while it seems to be able to sometimes pass through an epoch that it previously failed, it is not doing well, and once it hits nan it generally returns nan moving forward.
I managed to capture the error with detect_anomaly this time though, but unfortunately it doesn’t give me much to go on:

[W ..\torch\csrc\autograd\python_anomaly_mode.cpp:104] Warning: Error detected in NormBackward1. Traceback of forward call that caused the error:
  File "C:\Program Files\JetBrains\PyCharm Community Edition 2021.1.2\plugins\python-ce\helpers\pydev\pydevd.py", line 2173, in <module>
    main()
  File "C:\Program Files\JetBrains\PyCharm Community Edition 2021.1.2\plugins\python-ce\helpers\pydev\pydevd.py", line 2164, in main
    globals = debugger.run(setup['file'], None, None, is_module)
  File "C:\Program Files\JetBrains\PyCharm Community Edition 2021.1.2\plugins\python-ce\helpers\pydev\pydevd.py", line 1476, in run
    return self._exec(is_module, entry_point_fn, module_name, file, globals, locals)
  File "C:\Program Files\JetBrains\PyCharm Community Edition 2021.1.2\plugins\python-ce\helpers\pydev\pydevd.py", line 1483, in _exec
    pydev_imports.execfile(file, globals, locals)  # execute the script
  File "C:\Program Files\JetBrains\PyCharm Community Edition 2021.1.2\plugins\python-ce\helpers\pydev\_pydev_imps\_pydev_execfile.py", line 18, in execfile
    exec(compile(contents+"\n", file, 'exec'), glob, loc)
  File "C:/Users/Tue/PycharmProjects/ConstrainedMD/src/run_MD_e3_batch.py", line 99, in <module>
    results = main(c)
  File "C:\Users\Tue\PycharmProjects\ConstrainedMD\src\main.py", line 195, in main
    aloss_t, alossr_t, alossv_t, alossD_t, alossDr_t, alossDv_t, ap_t, MAEr_t, MAEv_t = run_network_e3(model, dataloader_train, train=True, max_samples=1e6, optimizer=optimizer, loss_fnc=c['loss'], batch_size=c['batch_size'], max_radius=cn['max_radius'], debug=c['debug'], log=LOG)
  File "C:\Users\Tue\PycharmProjects\ConstrainedMD\src\utils.py", line 195, in run_network_e3
    loss_v = torch.sum(torch.norm(Vpred - Vout_vec, p=2, dim=1)) / nb
  File "C:\Users\Tue\PycharmProjects\ConstrainedMD\venv\PycharmProjects\lib\site-packages\torch\functional.py", line 1420, in norm
    return _VF.norm(input, p, _dim, keepdim=keepdim)  # type: ignore
 (function _print_stack)
Traceback (most recent call last):
  File "C:\Program Files\JetBrains\PyCharm Community Edition 2021.1.2\plugins\python-ce\helpers\pydev\pydevd.py", line 1483, in _exec
    pydev_imports.execfile(file, globals, locals)  # execute the script
  File "C:\Program Files\JetBrains\PyCharm Community Edition 2021.1.2\plugins\python-ce\helpers\pydev\_pydev_imps\_pydev_execfile.py", line 18, in execfile
    exec(compile(contents+"\n", file, 'exec'), glob, loc)
  File "C:/Users/Tue/PycharmProjects/ConstrainedMD/src/run_MD_e3_batch.py", line 99, in <module>
    results = main(c)
  File "C:\Users\Tue\PycharmProjects\ConstrainedMD\src\main.py", line 195, in main
    aloss_t, alossr_t, alossv_t, alossD_t, alossDr_t, alossDv_t, ap_t, MAEr_t, MAEv_t = run_network_e3(model, dataloader_train, train=True, max_samples=1e6, optimizer=optimizer, loss_fnc=c['loss'], batch_size=c['batch_size'], max_radius=cn['max_radius'], debug=c['debug'], log=LOG)
  File "C:\Users\Tue\PycharmProjects\ConstrainedMD\src\utils.py", line 237, in run_network_e3
    loss.backward()
  File "C:\Users\Tue\PycharmProjects\ConstrainedMD\venv\PycharmProjects\lib\site-packages\torch\tensor.py", line 245, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)
  File "C:\Users\Tue\PycharmProjects\ConstrainedMD\venv\PycharmProjects\lib\site-packages\torch\autograd\__init__.py", line 145, in backward
    Variable._execution_engine.run_backward(
RuntimeError: Function 'NormBackward1' returned nan values in its 0th output.

I have been trying to figure out what is meant by its 0th output, but haven’t found much info anywhere on detect_anomaly.