[fx2trt][easy] add sparse flag to TRTInterpreter (pytorch#69495)

Summary:
Pull Request resolved: pytorch#69495

As the title. Separated from D30589161.

Test Plan: Tested in D30589161.

Reviewed By: maratsubkhankulov, wushirong

Differential Revision: D32898927

fbshipit-source-id: 89e18d2eb19b43fbab92b4988d0a21d21cff2d1f

torch/fx/experimental/fx2trt/fx2trt.py

@@ -143,6 +143,7 @@ def run(
         max_workspace_size=1 << 25,
         fp16_mode=True,
         int8_mode=False,
+        sparse_weights=False,
         force_fp32_output=False,
         strict_type_constraints=False,
         algorithm_selector=None,
@@ -179,6 +180,10 @@ def run(
         if int8_mode:
             builder_config.set_flag(trt.BuilderFlag.INT8)
 
+        if sparse_weights:
+            assert fp16_mode or int8_mode, "We can only enable sparsity in fp16 or int8 mode."
+            builder_config.set_flag(trt.BuilderFlag.SPARSE_WEIGHTS)
+
         if strict_type_constraints:
             builder_config.set_flag(trt.BuilderFlag.STRICT_TYPES)