test

Author: lanluo-nvidia

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

@@ -661,16 +661,7 @@ def aten_ops_unsqueeze(
     kwargs: Dict[str, Argument],
     name: str,
 ) -> Union[TRTTensor, Sequence[TRTTensor]]:
-    from importlib.metadata import version
-
-    if version("tensorrt") >= "10.7.0":
-        return impl.unsqueeze.unsqueeze(
-            ctx, target, SourceIR.ATEN, name, args[0], args[1]
-        )
-    else:
-        return impl.unsqueeze.unsqueeze_old(
-            ctx, target, SourceIR.ATEN, name, args[0], args[1]
-        )
+    return impl.unsqueeze.unsqueeze(ctx, target, SourceIR.ATEN, name, args[0], args[1])
 
 
 @dynamo_tensorrt_converter(

py/torch_tensorrt/dynamo/conversion/impl/unsqueeze.py

@@ -19,6 +19,10 @@ def unsqueeze(
     input: TRTTensor,
     dim: int,
 ) -> TRTTensor:
+    from importlib.metadata import version
+
+    if version("tensorrt") < "10.7.0":
+        return unsqueeze_old(ctx, target, source_ir, name, input, dim)
     axes = get_trt_tensor(ctx, dim, f"{name}_axes")
     layer = ctx.net.add_unsqueeze(input, axes)
     set_layer_name(layer, target, name, source_ir)