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 1
1
import os
2
+ import re
2
3
from io import BytesIO
3
4
from pathlib import Path
4
5
5
6
import pytest
6
7
import torch
7
8
8
9
import tests_pytorch .helpers .pipelines as pipes
10
+ from lightning .pytorch .core .module import _TORCH_TRT_AVAILABLE
9
11
from lightning .pytorch .demos .boring_classes import BoringModel
10
12
from tests_pytorch .helpers .runif import RunIf
11
13
12
14
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
+
13
25
@RunIf (tensorrt = True , min_cuda_gpus = 1 , min_torch = "2.2.0" )
14
26
def test_tensorrt_saves_with_input_sample (tmp_path ):
15
27
model = BoringModel ()
You can’t perform that action at this time.
0 commit comments