deeplearning/projects/flashlight-sequence

Reviewed By: zsol

Differential Revision: D53083529

fbshipit-source-id: d834e9535aba755da0df1167ffb19b59ca8e9933

Author: generatedunixname89002005287564

bindings/python/example/criterion_example.py

@@ -12,7 +12,7 @@
 from flashlight.lib.sequence.criterion import ASGLoss, CriterionScaleMode
 
 
-if __name__ == "__main__":
+def main() -> None:
     parser = argparse.ArgumentParser()
     parser.add_argument(
         "--cpu", action="store_true", help="Use cpu backend, otherwise use CUDA backend"
@@ -80,3 +80,7 @@
     loss.backward(grad)
     print("Gradients with respect to the ASG loss input", input.grad)
     print("Gradients with respect to the transition matrix", asg.transitions.grad)
+
+
+if __name__ == "__main__":
+    main()  # pragma: no cover