Improve checkpoint handling for neural networks

Fixes #56