setup amp correctly

Author: lucidrains

rin_pytorch/rin_pytorch.py

@@ -906,20 +906,18 @@ def __init__(
         num_samples = 25,
         results_folder = './results',
         amp = False,
-        fp16 = False,
+        mixed_precision_type = 'fp16',
         split_batches = True,
         convert_image_to = None
     ):
         super().__init__()
 
         self.accelerator = Accelerator(
             split_batches = split_batches,
-            mixed_precision = 'fp16' if fp16 else 'no',
+            mixed_precision = mixed_precision_type if amp else 'no',
             kwargs_handlers = [DistributedDataParallelKwargs(find_unused_parameters=True)]
         )
 
-        self.accelerator.native_amp = amp
-
         self.model = diffusion_model
 
         assert has_int_squareroot(num_samples), 'number of samples must have an integer square root'

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'RIN-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.7.5',
+  version = '0.7.6',
   license='MIT',
   description = 'RIN - Recurrent Interface Network - Pytorch',
   author = 'Phil Wang',