diff --git a/docker/Dockerfile-jetson-jetpack5 b/docker/Dockerfile-jetson-jetpack5
index 1a0013d8d2..3ac12963fa 100644
--- a/docker/Dockerfile-jetson-jetpack5
+++ b/docker/Dockerfile-jetson-jetpack5
@@ -41,8 +41,8 @@ RUN sed -i 's/^\( *"tensorflowjs\)>=.*\(".*\)/\1>=3.9.0\2/' pyproject.toml && \
 # Pip install onnxruntime-gpu, torch, torchvision and ultralytics, then remove build files
 RUN uv pip install --system \
     https://github.com/ultralytics/assets/releases/download/v0.0.0/onnxruntime_gpu-1.18.0-cp38-cp38-linux_aarch64.whl \
-    https://github.com/ultralytics/assets/releases/download/v0.0.0/torch-2.2.0-cp38-cp38-linux_aarch64.whl \
-    https://github.com/ultralytics/assets/releases/download/v0.0.0/torchvision-0.17.2+c1d70fe-cp38-cp38-linux_aarch64.whl && \
+    https://github.com/ultralytics/assets/releases/download/v0.0.0/torch-2.1.0a0+41361538.nv23.06-cp38-cp38-linux_aarch64.whl \
+    https://github.com/ultralytics/assets/releases/download/v0.0.0/torchvision-0.16.2+c6f3977-cp38-cp38-linux_aarch64.whl && \
     # Need lower version of 'numpy' for TensorRT export
     uv pip install --system numpy==1.23.5 && \
     uv pip install --system -e ".[export]" && \
diff --git a/docs/en/guides/nvidia-jetson.md b/docs/en/guides/nvidia-jetson.md
index fc066bef54..f380069ca6 100644
--- a/docs/en/guides/nvidia-jetson.md
+++ b/docs/en/guides/nvidia-jetson.md
@@ -270,11 +270,11 @@ The above ultralytics installation will install Torch and Torchvision. However,
     pip uninstall torch torchvision
     ```
 
-2. Install `torch 2.2.0` and `torchvision 0.17.2` according to JP5.1.2
+2. Install `torch 2.1.0` and `torchvision 0.16.2` according to JP5.1.2
 
     ```bash
-    pip install https://github.com/ultralytics/assets/releases/download/v0.0.0/torch-2.2.0-cp38-cp38-linux_aarch64.whl
-    pip install https://github.com/ultralytics/assets/releases/download/v0.0.0/torchvision-0.17.2+c1d70fe-cp38-cp38-linux_aarch64.whl
+    pip install https://github.com/ultralytics/assets/releases/download/v0.0.0/torch-2.1.0a0+41361538.nv23.06-cp38-cp38-linux_aarch64.whl
+    pip install https://github.com/ultralytics/assets/releases/download/v0.0.0/torchvision-0.16.2+c6f3977-cp38-cp38-linux_aarch64.whl
     ```
 
 !!! note
diff --git a/ultralytics/engine/trainer.py b/ultralytics/engine/trainer.py
index 49b9d937ac..05b6e43bd1 100644
--- a/ultralytics/engine/trainer.py
+++ b/ultralytics/engine/trainer.py
@@ -631,10 +631,7 @@ class BaseTrainer:
         import io
 
         ema = deepcopy(unwrap_model(self.ema.ema)).half()
-        if (
-            not all(torch.isfinite(v).all() for v in ema.state_dict().values() if isinstance(v, torch.Tensor))
-            and self.epoch > self.start_epoch  # at least save checkpoint for the first epoch
-        ):
+        if not all(torch.isfinite(v).all() for v in ema.state_dict().values() if isinstance(v, torch.Tensor)):
             LOGGER.warning(f"Skipping checkpoint save at epoch {self.epoch}: EMA contains NaN/Inf")
             return False