Merge pull request #182 from cuuupid/patch-1

Declare `amp_backend` parameter for Torch Lightning
minimaxir · May 23, 2022 · 29d8dfa · 29d8dfa
2 parents 24d9291 + e9f66ba
commit 29d8dfa
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/aitextgen/aitextgen.py b/aitextgen/aitextgen.py
@@ -734,6 +734,7 @@ def train(
         if fp16:
             train_params["precision"] = 16 if fp16 else 32
             train_params["amp_level"] = fp16_opt_level
+            train_params["amp_backend"] = "apex"
 
         if tpu_cores > 0:
             train_params["tpu_cores"] = tpu_cores