skipping test_refit_weight_stripped_engine_multiple_times if include_refit flag not there

Author: apbose

tests/py/dynamo/models/test_weight_stripped_engine.py

@@ -4,6 +4,7 @@
 import shutil
 import unittest
 
+import tensorrt as trt
 import torch
 import torch_tensorrt as torch_trt
 from torch.testing._internal.common_utils import TestCase
@@ -642,6 +643,10 @@ def test_refit_identical_engine_weights(self):
         not importlib.util.find_spec("torchvision"),
         "torchvision is not installed",
     )
+    @unittest.skipIf(
+        not hasattr(trt.SerializationFlag, "INCLUDE_REFIT"),
+        "Multiple refit requires TensorRT >= 10.14 with INCLUDE_REFIT serialization flag",
+    )
     def test_refit_weight_stripped_engine_multiple_times(self):
         pyt_model = models.resnet18(pretrained=True).eval().to("cuda")
         example_inputs = (torch.randn((100, 3, 224, 224)).to("cuda"),)

tests/py/ts/api/test_classes.py

@@ -10,20 +10,22 @@
 def is_blackwell():
     """
     Check if running on NVIDIA Blackwell architecture (sm_90+).
-    
+
     Blackwell architecture adds input/output reformat layers in TensorRT engines.
-    
+
     Returns:
         bool: True if running on Blackwell (sm_90+), False otherwise
     """
     if not torch.cuda.is_available():
         return False
-    
+
     device_properties = torch.cuda.get_device_properties(0)
     compute_capability = device_properties.major * 10 + device_properties.minor
-    
+
     # Blackwell is sm_90 and above
     return compute_capability >= 90
+
+
 @unittest.skipIf(
     not torchtrt.ENABLED_FEATURES.torchscript_frontend,
     "TorchScript Frontend is not available",
@@ -348,12 +350,13 @@ def test_get_layer_info(self):
         """
 
         import json
+
         if is_blackwell():
-            # blackwell has additional layers- 
-            #Layer 0: __mye88_myl0_0           ← Input reformat layer
-            #Layer 1: aten__matmul(...) fc1    ← First matmul (fc1)
-            #Layer 2: aten__matmul(...) fc2    ← Second matmul (fc2)
-            #Layer 3: __mye90_myl0_3           ← Output reformat layer
+            # blackwell has additional layers-
+            # Layer 0: __mye88_myl0_0           ← Input reformat layer
+            # Layer 1: aten__matmul(...) fc1    ← First matmul (fc1)
+            # Layer 2: aten__matmul(...) fc2    ← Second matmul (fc2)
+            # Layer 3: __mye90_myl0_3           ← Output reformat layer
             num_layers = 4
         else:
             num_layers = 2