diff --git a/torchrec_dlrm/dlrm_main.py b/torchrec_dlrm/dlrm_main.py
index 99f95d01..5b02bbbd 100644
--- a/torchrec_dlrm/dlrm_main.py
+++ b/torchrec_dlrm/dlrm_main.py
@@ -719,5 +719,9 @@ def optimizer_with_params():
         multihot.save_freqs_stats()
 
 
-if __name__ == "__main__":
+def invoke_main() -> None:
     main(sys.argv[1:])
+
+
+if __name__ == "__main__":
+    invoke_main()  # pragma: no cover