Regression: MNIST example produces NaN for loss

[antimora]

Describe the bug
Running the training gives NaN for loss:

[Metrics]
  - Train Accuracy: epoch 0.00 % - batch 0.00 %
  - Train Loss: epoch NaN - batch NaN
  - Valid Accuracy: epoch 0.00 % - batch 0.00 %
  - Valid Loss: epoch NaN - batch NaN

[Progress]
  - Iteration 1422 Epoch 2/6
  - iteration [############################>-----------------------------------------------------------------------] (84s)
  - epoch     [#################################>-------------------------------------------------------------------] (4m)

To Reproduce

1. Check out the latest commit: be96160
2. cd to examples/mnist
3. run: cargo run --example mnist --release --features tch-gpu

Expected behavior

Screenshots

Desktop (please complete the following information):

[mnist]$ uname -a
Darwin MacBook-Pro.local 22.3.0 Darwin Kernel Version 22.3.0: 
Mon Jan 30 20:38:37 PST 2023; root:xnu-8792.81.3~2/RELEASE_ARM64_T6000 arm64

Additional context
It worked on a commit before a few days.