chore: fix docs for export

Signed-off-by: Dheeraj Peri <[email protected]>

chore: updates

Signed-off-by: Dheeraj Peri <[email protected]>

chore: updates

Signed-off-by: Dheeraj Peri <[email protected]>

Author: peri044

docsrc/dynamo/dynamo_export.rst

@@ -1,6 +1,6 @@
 .. _dynamo_export:
 
-Compiling ``ExportedPrograms`` with Torch-TensorRT
+Compiling Exported Programs with Torch-TensorRT
 =============================================
 .. currentmodule:: torch_tensorrt.dynamo
 
@@ -9,8 +9,6 @@ Compiling ``ExportedPrograms`` with Torch-TensorRT
    :undoc-members:
    :show-inheritance:
 
-Using the Torch-TensorRT Frontend for ``torch.export.ExportedPrograms``
---------------------------------------------------------
 Pytorch 2.1 introduced ``torch.export`` APIs which
 can export graphs from Pytorch programs into ``ExportedProgram`` objects. Torch-TensorRT dynamo
 frontend compiles these ``ExportedProgram`` objects and optimizes them using TensorRT. Here's a simple

docsrc/user_guide/saving_models.rst

@@ -29,14 +29,14 @@ The following code illustrates this approach.
     import torch_tensorrt
 
     model = MyModel().eval().cuda()
-    inputs = torch.randn((1, 3, 224, 224)).cuda()
+    inputs = [torch.randn((1, 3, 224, 224)).cuda()]
     trt_gm = torch_tensorrt.compile(model, ir="dynamo", inputs) # Output is a torch.fx.GraphModule
-    trt_script_model = torch.jit.trace(trt_gm, inputs)
-    torch.jit.save(trt_script_model, "trt_model.ts")
+    trt_traced_model = torch.jit.trace(trt_gm, inputs)
+    torch.jit.save(trt_traced_model, "trt_model.ts")
 
     # Later, you can load it and run inference
     model = torch.jit.load("trt_model.ts").cuda()
-    model(inputs)
+    model(*inputs)
 
 b) ExportedProgram
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
@@ -50,16 +50,15 @@ b) ExportedProgram
     import torch_tensorrt
 
     model = MyModel().eval().cuda()
-    inputs = torch.randn((1, 3, 224, 224)).cuda()
+    inputs = [torch.randn((1, 3, 224, 224)).cuda()]
     trt_gm = torch_tensorrt.compile(model, ir="dynamo", inputs) # Output is a torch.fx.GraphModule
     # Transform and create an exported program
-    trt_gm = torch_tensorrt.dynamo.export(trt_gm, inputs)
-    trt_exp_program = create_exported_program(trt_gm, call_spec, trt_gm.state_dict())
-    torch._export.save(trt_exp_program, "trt_model.ep")
+    trt_exp_program = torch_tensorrt.dynamo.export(trt_gm, inputs)
+    torch.export.save(trt_exp_program, "trt_model.ep")
 
     # Later, you can load it and run inference
-    model = torch._export.load("trt_model.ep")
-    model(inputs)
+    model = torch.export.load("trt_model.ep")
+    model(*inputs)
 
 `torch_tensorrt.dynamo.export` inlines the submodules within a GraphModule to their corresponding nodes and stiches all the nodes together.
 This is needed as `torch._export` serialization cannot handle serializing and deserializing of submodules (`call_module` nodes).
@@ -79,11 +78,11 @@ Torchscript IR
     import torch_tensorrt
 
     model = MyModel().eval().cuda()
-    inputs = torch.randn((1, 3, 224, 224)).cuda()
+    inputs = [torch.randn((1, 3, 224, 224)).cuda()]
     trt_ts = torch_tensorrt.compile(model, ir="ts", inputs) # Output is a ScriptModule object
     torch.jit.save(trt_ts, "trt_model.ts")
 
     # Later, you can load it and run inference
     model = torch.jit.load("trt_model.ts").cuda()
-    model(inputs)
+    model(*inputs)
 

py/torch_tensorrt/dynamo/_exporter.py

@@ -56,7 +56,7 @@ def export(
         return exp_program
     else:
         raise ValueError(
-            "Invalid ir : {ir} provided for serialization. Options include torchscript | exported_program"
+            f"Invalid ir : {ir} provided for serialization. Options include torchscript | exported_program"
         )