fix amp

lucidrains · Sep 4, 2023 · bac049a · bac049a
1 parent deb97db
commit bac049a
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 5 deletions.
diff --git a/naturalspeech2_pytorch/naturalspeech2_pytorch.py b/naturalspeech2_pytorch/naturalspeech2_pytorch.py
@@ -1649,7 +1649,7 @@ def __init__(
         num_samples = 1,
         results_folder = './results',
         amp = False,
-        fp16 = False,
+        mixed_precision_type = 'fp16',
         use_ema = True,
         split_batches = True,
         dataloader = None,
@@ -1663,11 +1663,9 @@ def __init__(
 
         self.accelerator = Accelerator(
             split_batches = split_batches,
-            mixed_precision = 'fp16' if fp16 else 'no'
+            mixed_precision = mixed_precision_type if amp else 'no'
         )
 
-        self.accelerator.native_amp = amp
-
         # model
 
         self.model = diffusion_model

diff --git a/naturalspeech2_pytorch/version.py b/naturalspeech2_pytorch/version.py
@@ -1 +1 @@
-__version__ = '0.1.4'
+__version__ = '0.1.5'