Skip to content

Commit 484f9ce

Browse files
committed
test: add test test_missing_tensorrt_package.
1 parent c97ae0d commit 484f9ce

File tree

1 file changed

+12
-0
lines changed

1 file changed

+12
-0
lines changed

tests/tests_pytorch/models/test_torch_tensorrt.py

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,27 @@
11
import os
2+
import re
23
from io import BytesIO
34
from pathlib import Path
45

56
import pytest
67
import torch
78

89
import tests_pytorch.helpers.pipelines as pipes
10+
from lightning.pytorch.core.module import _TORCH_TRT_AVAILABLE
911
from lightning.pytorch.demos.boring_classes import BoringModel
1012
from tests_pytorch.helpers.runif import RunIf
1113

1214

15+
@pytest.mark.skipif(_TORCH_TRT_AVAILABLE, reason="Run this test only if tensorrt is not available.")
16+
def test_missing_tensorrt_package():
17+
model = BoringModel()
18+
with pytest.raises(
19+
ModuleNotFoundError,
20+
match=re.escape(f"`{type(model).__name__}.to_tensorrt` requires `torch_tensorrt` to be installed. "),
21+
):
22+
model.to_tensorrt("model.trt")
23+
24+
1325
@RunIf(tensorrt=True, min_cuda_gpus=1, min_torch="2.2.0")
1426
def test_tensorrt_saves_with_input_sample(tmp_path):
1527
model = BoringModel()

0 commit comments

Comments
 (0)