Merge pull request #64 from taiypeo/feature/amp_backend

Add support for different AMP backends
HHousen · Feb 2, 2022 · e39140c · e39140c
2 parents 1395f37 + 7ffef86
commit e39140c
Showing 1 changed file with 8 additions and 1 deletion.
diff --git a/src/main.py b/src/main.py
@@ -241,10 +241,17 @@ def main(args):
     parser.add_argument(
         "--amp_level",
         type=str,
-        default="O1",
+        default=None,
         help="The optimization level to use (O1, O2, etc…) for 16-bit GPU precision (using "
         + "NVIDIA apex under the hood).",
     )
+    parser.add_argument(
+        "--amp_backend",
+        type=str,
+        default="native",
+        choices=["native", "apex"],
+        help="PyTorch Lightning amp_backend ('native' or 'apex')",
+    )
     parser.add_argument(
         "--precision",
         type=int,