Elgato_dark/ultralytics
1
0
Fork 0
mirror of https://github.com/ultralytics/ultralytics synced 2026-04-21 14:07:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 8

Fix torch and torchvision versions on JetPack 5 Docker (#24199)

Browse source
This commit is contained in:
Lakshantha Dissanayake 2026-04-10 18:37:33 -07:00 committed by GitHub
parent f06a9bbadd
commit c7e1d6044a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 6 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
docker/Dockerfile-jetson-jetpack5
View file

@ -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]" && \

6
docs/en/guides/nvidia-jetson.md
View file

@ -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

5
ultralytics/engine/trainer.py
View file

@ -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