Fix forward unittests.

Summary: from D49482310

Reviewed By: JacobSzwejbka, SherlockNoMad

Differential Revision: D49547858

fbshipit-source-id: fb4c3b397174c585394511d7baf097fd7023c441

Author: zhxchen17

exir/program/test/test_program.py

@@ -156,7 +156,7 @@ def test_edge_manager_transform(self):
         )
 
         original_res = edge_manager.exported_program("forward")(
-            torch.ones(1), torch.ones(1), torch.ones(1)
+            torch.ones(1), torch.ones(1)
         )
 
         # perform transformation
@@ -173,17 +173,13 @@ def test_edge_manager_transform(self):
 
         # transformation was applied
         self.assertEqual(
-            transformed_edge.exported_program("forward")(
-                torch.ones(1), torch.ones(1), torch.ones(1)
-            ),
+            transformed_edge.exported_program("forward")(torch.ones(1), torch.ones(1)),
             torch.ones(1),  # x * y * x
         )
 
         # original unchanged
         self.assertEqual(
-            edge_manager.exported_program("forward")(
-                torch.ones(1), torch.ones(1), torch.ones(1)
-            ),
+            edge_manager.exported_program("forward")(torch.ones(1), torch.ones(1)),
             original_res,  # x * y + x
         )
 
@@ -199,9 +195,7 @@ def test_transform_dict_api(self):
         )
 
         self.assertEqual(
-            transformed_edge.exported_program("forward")(
-                torch.ones(1), torch.ones(1), torch.ones(1)
-            ),
+            transformed_edge.exported_program("forward")(torch.ones(1), torch.ones(1)),
             torch.ones(1),  # x * y * x
         )
 
@@ -222,7 +216,7 @@ def test_edge_to_backend_replaces_subgraph(self):
 
         forward_program = delegate_manager.exported_program("forward")
         self.assertEqual(
-            forward_program(torch.ones(1), torch.ones(1), torch.ones(1)),
+            forward_program(torch.ones(1), torch.ones(1)),
             torch.ones(1) + 1,  # x * y + x
         )