torchtnt

Differential Revision: D51846380
pytorch · Dec 5, 2023 · 535654d · 535654d
1 parent 80b4b0c
commit 535654d
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 4 deletions.
diff --git a/examples/auto_unit_example.py b/examples/auto_unit_example.py
@@ -239,7 +239,7 @@ def get_args() -> Namespace:
     return parser.parse_args()
 
 
-if __name__ == "__main__":
+def invoke_main() -> None:
     args: Namespace = get_args()
     lc = pet.LaunchConfig(
         min_nodes=1,
@@ -253,3 +253,7 @@ def get_args() -> Namespace:
     )
 
     pet.elastic_launch(lc, entrypoint=main)(args)
+
+
+if __name__ == "__main__":
+    invoke_main()  # pragma: no cover
diff --git a/examples/torchdata_train_example.py b/examples/torchdata_train_example.py
@@ -202,5 +202,9 @@ def get_args(argv: List[str]) -> Namespace:
     return parser.parse_args(argv)
 
 
-if __name__ == "__main__":
+def invoke_main() -> None:
     main(sys.argv[1:])
+
+
+if __name__ == "__main__":
+    invoke_main()  # pragma: no cover
diff --git a/examples/torchrec/main.py b/examples/torchrec/main.py
@@ -378,7 +378,7 @@ def main(argv: List[str]) -> None:
     )
 
 
-if __name__ == "__main__":
+def invoke_main() -> None:
     lc = launcher.LaunchConfig(
         min_nodes=MIN_NODES,
         max_nodes=MAX_NODES,
@@ -391,3 +391,7 @@ def main(argv: List[str]) -> None:
     )
 
     launcher.elastic_launch(config=lc, entrypoint=main)(sys.argv[1:])
+
+
+if __name__ == "__main__":
+    invoke_main()  # pragma: no cover
diff --git a/examples/train_unit_example.py b/examples/train_unit_example.py
@@ -173,5 +173,9 @@ def get_args(argv: List[str]) -> Namespace:
     return parser.parse_args(argv)
 
 
-if __name__ == "__main__":
+def invoke_main() -> None:
     main(sys.argv[1:])
+
+
+if __name__ == "__main__":
+    invoke_main()  # pragma: no cover