Update optimization.py

optimization.py

@@ -42,14 +42,14 @@ def optimize_pipeline_(pipeline: Callable[P, Any], *args: P.args, **kwargs: P.kw
         
         exported_A = torch.export.export(
             mod=blocks_A[0],
-            args=call.args,
-            kwargs=call.kwargs,
+            args=call_A.args,
+            kwargs=call_A.kwargs,
         )
 
         exported_B = torch.export.export(
             mod=blocks_B[0],
-            args=call.args,
-            kwargs=call.kwargs,
+            args=call_B.args,
+            kwargs=call_B.kwargs,
         )
 
         return (