Skip to content

Commit

Permalink
Skip test without import
Browse files Browse the repository at this point in the history
  • Loading branch information
meatybobby committed Nov 26, 2024
1 parent eda258c commit 27af4a2
Showing 1 changed file with 0 additions and 8 deletions.
8 changes: 0 additions & 8 deletions tests/export/test_trt_compile.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,11 +20,6 @@
from nemo.export import trt_compile
from nemo.utils.import_utils import safe_import

trt, trt_imported = safe_import("tensorrt")
torch_tensorrt, torch_trt_imported = safe_import("torch_tensorrt")
polygraphy, polygraphy_imported = safe_import("polygraphy")


TEST_CASE_1 = ["fp32"]
TEST_CASE_2 = ["fp16"]

Expand All @@ -43,8 +38,6 @@ def forward(self, x: List[torch.Tensor], y: torch.Tensor, z: torch.Tensor, bs: f


@unittest.skip
@unittest.skipUnless(trt_imported, "tensorrt is required")
@unittest.skipUnless(polygraphy_imported, "polygraphy is required")
class TestTRTCompile(unittest.TestCase):

def setUp(self):
Expand All @@ -55,7 +48,6 @@ def tearDown(self):
if current_device != self.gpu_device:
torch.cuda.set_device(self.gpu_device)

@unittest.skipUnless(torch_trt_imported, "torch_tensorrt is required")
def test_torch_trt(self):

model = torch.nn.Sequential(*[torch.nn.PReLU(), torch.nn.PReLU()])
Expand Down

0 comments on commit 27af4a2

Please sign in to comment.