diff --git a/py/torch_tensorrt/dynamo/conversion/custom_ops_converters.py b/py/torch_tensorrt/dynamo/conversion/custom_ops_converters.py index 7bd70705a0..36386b8d7b 100644 --- a/py/torch_tensorrt/dynamo/conversion/custom_ops_converters.py +++ b/py/torch_tensorrt/dynamo/conversion/custom_ops_converters.py @@ -28,7 +28,9 @@ _LOGGER.info("Using native TensorRT DistCollective API for distributed operations") @dynamo_tensorrt_converter( - tensorrt_fused_nccl_all_gather_op, requires_native_multidevice=True + tensorrt_fused_nccl_all_gather_op, + requires_native_multidevice=True, + supports_dynamic_shapes=True, ) def fused_nccl_all_gather( ctx: ConversionContext, @@ -47,7 +49,9 @@ def fused_nccl_all_gather( ) @dynamo_tensorrt_converter( - tensorrt_fused_nccl_reduce_scatter_op, requires_native_multidevice=True + tensorrt_fused_nccl_reduce_scatter_op, + requires_native_multidevice=True, + supports_dynamic_shapes=True, ) def fused_nccl_reduce_scatter( ctx: ConversionContext, @@ -69,7 +73,9 @@ def fused_nccl_reduce_scatter( ) @dynamo_tensorrt_converter( - tensorrt_fused_nccl_all_reduce_op, requires_native_multidevice=True + tensorrt_fused_nccl_all_reduce_op, + requires_native_multidevice=True, + supports_dynamic_shapes=True, ) def fused_nccl_all_reduce( ctx: ConversionContext, @@ -90,7 +96,9 @@ def fused_nccl_all_reduce( ) @dynamo_tensorrt_converter( - tensorrt_fused_nccl_all_to_all_op, requires_native_multidevice=True + tensorrt_fused_nccl_all_to_all_op, + requires_native_multidevice=True, + supports_dynamic_shapes=True, ) def fused_nccl_all_to_all( ctx: ConversionContext, @@ -109,7 +117,9 @@ def fused_nccl_all_to_all( ) @dynamo_tensorrt_converter( - tensorrt_fused_nccl_scatter_op, requires_native_multidevice=True + tensorrt_fused_nccl_scatter_op, + requires_native_multidevice=True, + supports_dynamic_shapes=True, ) def fused_nccl_scatter( ctx: ConversionContext, @@ -125,7 +135,9 @@ def fused_nccl_scatter( ) @dynamo_tensorrt_converter( - tensorrt_fused_nccl_gather_op, requires_native_multidevice=True + tensorrt_fused_nccl_gather_op, + requires_native_multidevice=True, + supports_dynamic_shapes=True, ) def fused_nccl_gather( ctx: ConversionContext,