utility function to detect tegra platform

Author: apbose

py/torch_tensorrt/dynamo/lowering/passes/_aten_lowering_pass.py

@@ -3,6 +3,7 @@
 
 import torch
 from torch_tensorrt.dynamo._settings import CompilationSettings
+from torch_tensorrt.dynamo.utils import is_tegra_platform
 
 from .accumulate_fp32_matmul import accumulate_fp32_matmul
 from .constant_folding import constant_fold
@@ -25,11 +26,11 @@
     replace_max_pool_with_indices,
     lower_scaled_dot_product_attention,
     view_to_reshape,
-    remove_assert_nodes,
+    remove_assert_scalar,
     accumulate_fp32_matmul,
 ]
 
-if torch.cuda.get_device_capability() not in [(8, 7), (7, 2)]:
+if not is_tegra_platform():
     pass_list.append(fuse_distributed_ops)
 
 ATEN_POST_LOWERING_PASSES = DynamoPassManager.build_from_passlist(pass_list)

py/torch_tensorrt/dynamo/utils.py

@@ -793,3 +793,9 @@ def get_output_dtypes(output: Any, truncate_doulbe: bool = False) -> List[dtype]
             f"got unexpected type {type(output)}, expected type is a torch.fx.node.Node or a tuple/list of torch.fx.node.Node"
         )
     return output_dtypes
+
+
+def is_tegra_platform() -> bool:
+    if torch.cuda.get_device_capability() in [(8, 7), (7, 2)]:
+        return True
+    return False