File tree Expand file tree Collapse file tree 1 file changed +12
-0
lines changed
tests/tests_pytorch/models Expand file tree Collapse file tree 1 file changed +12
-0
lines changed Original file line number Diff line number Diff line change 11import os
2+ import re
23from io import BytesIO
34from pathlib import Path
45
56import pytest
67import torch
78
89import tests_pytorch .helpers .pipelines as pipes
10+ from lightning .pytorch .core .module import _TORCH_TRT_AVAILABLE
911from lightning .pytorch .demos .boring_classes import BoringModel
1012from 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" )
1426def test_tensorrt_saves_with_input_sample (tmp_path ):
1527 model = BoringModel ()
You can’t perform that action at this time.
0 commit comments