HongChuang/ultralytics-ascend
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 3

ultralytics 8.1.46 add TensorRT 10 support (#9516)

Browse source 
Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
Co-authored-by: 九是否随意的称呼 <1069679911@qq.com>
This commit is contained in:
Burhan 2024-04-10 12:07:53 -04:00 • committed by GitHub
parent ea03db9984
commit 4ffd6ee6d7
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 77 additions and 32 deletions
Download patch file Download diff file Expand all files Collapse all files

7
tests/test_cuda.py
View file

@ -19,6 +19,13 @@ def test_checks():
assert torch.cuda.is_available() == CUDA_IS_AVAILABLE
assert torch.cuda.device_count() == CUDA_DEVICE_COUNT
@pytest.mark.slow
@pytest.mark.skipif(not CUDA_IS_AVAILABLE, reason="CUDA is not available")
def test_export_engine():
"""Test exporting the YOLO model to NVIDIA TensorRT format."""
f = YOLO(MODEL).export(format="engine", device=0)
YOLO(f)(BUS, device=0)
@pytest.mark.skipif(not CUDA_IS_AVAILABLE, reason="CUDA is not available")
def test_train():