Issue with torchvision when training YOLO model on Jetson

Hello. I am trying to train a yolov8 model on the Jetson Orin AGX, but I keep running into this error:

Traceback (most recent call last):
File “/home/buckeyevertical/Documents/training/train.py”, line 17, in
results = model.train(data = r"/home/buckeyevertical/Documents/training/data.yaml",
File “/home/buckeyevertical/.local/lib/python3.10/site-packages/ultralytics/engine/model.py”, line 338, in train
self.trainer.train()
File “/home/buckeyevertical/.local/lib/python3.10/site-packages/ultralytics/engine/trainer.py”, line 190, in train
self._do_train(world_size)
File “/home/buckeyevertical/.local/lib/python3.10/site-packages/ultralytics/engine/trainer.py”, line 290, in _do_train
self._setup_train(world_size)
File “/home/buckeyevertical/.local/lib/python3.10/site-packages/ultralytics/engine/trainer.py”, line 240, in _setup_train
self.amp = torch.tensor(check_amp(self.model), device=self.device)
File “/home/buckeyevertical/.local/lib/python3.10/site-packages/ultralytics/utils/checks.py”, line 607, in check_amp
assert amp_allclose(YOLO(‘yolov8n.pt’), im)
File “/home/buckeyevertical/.local/lib/python3.10/site-packages/ultralytics/utils/checks.py”, line 595, in amp_allclose
a = m(im, device=device, verbose=False)[0].boxes.data # FP32 inference
File “/home/buckeyevertical/.local/lib/python3.10/site-packages/ultralytics/engine/model.py”, line 98, in call
return self.predict(source, stream, **kwargs)
File “/home/buckeyevertical/.local/lib/python3.10/site-packages/ultralytics/engine/model.py”, line 239, in predict
return self.predictor.predict_cli(source=source) if is_cli else self.predictor(source=source, stream=stream)
File “/home/buckeyevertical/.local/lib/python3.10/site-packages/ultralytics/engine/predictor.py”, line 198, in call
return list(self.stream_inference(source, model, *args, **kwargs)) # merge list of Result into one
File “/home/buckeyevertical/.local/lib/python3.10/site-packages/torch/utils/_contextlib.py”, line 35, in generator_context
response = gen.send(None)
File “/home/buckeyevertical/.local/lib/python3.10/site-packages/ultralytics/engine/predictor.py”, line 269, in stream_inference
self.results = self.postprocess(preds, im, im0s)
File “/home/buckeyevertical/.local/lib/python3.10/site-packages/ultralytics/models/yolo/detect/predict.py”, line 25, in postprocess
preds = ops.non_max_suppression(preds,
File “/home/buckeyevertical/.local/lib/python3.10/site-packages/ultralytics/utils/ops.py”, line 238, in non_max_suppression
i = torchvision.ops.nms(boxes, scores, iou_thres) # NMS
File “/home/buckeyevertical/.local/lib/python3.10/site-packages/torchvision/ops/boxes.py”, line 40, in nms
_assert_has_ops()
File “/home/buckeyevertical/.local/lib/python3.10/site-packages/torchvision/extension.py”, line 46, in _assert_has_ops
raise RuntimeError(
RuntimeError: Couldn’t load custom C++ ops. This can happen if your PyTorch and torchvision versions are incompatible, or if you had errors while compiling torchvision from source. For further information on the compatible versions, check GitHub - pytorch/vision: Datasets, Transforms and Models specific to Computer Vision for the compatibility matrix. Please check your PyTorch version with torch.version and your torchvision version with torchvision.version and verify if they are compatible, and if not please reinstall torchvision so that it matches your PyTorch install.

I have built pytorch 2.1.0 and torchvision 0.16.1 from this announcement instructions:

Any help would be greatly appreciated.

Hi,

Could you help to verify the PyTorch and TorchVision version and share the output with us?

$ python3
>>> import torch
>>> torch.__version__
$ python3
>>> import torchvision
>>> torchvision.__version__

Thanks.

Just ran this. Looks like the torchvision version was switching branches to 0.16.2. I uninstalled and reinstalled, and it seems to be working now. Thank you!

This topic was automatically closed 14 days after the last reply. New replies are no longer allowed.