x-evolution 0.1.19__py3-none-any.whl → 0.1.23__py3-none-any.whl

x_evolution/x_evolution.py

@@ -112,6 +112,8 @@ class EvoStrategy(Module):
                 self.print('pre_main_callback detected on environment passed in and is invoked')
                 environment.pre_main_callback()
 
+        accelerator.wait_for_everyone()
+
         # take care of model and parameters
 
         if isinstance(model, list):
@@ -120,6 +122,15 @@ class EvoStrategy(Module):
         self.model = model
         self.noisable_model = Noisable(model, low_rank = noise_low_rank)
 
+        # use prepare and run through environment once to sync params
+
+        wrapped_model = accelerator.prepare(model)
+
+        with torch.no_grad():
+            environment(wrapped_model)
+
+        # get param dictionary
+
         named_parameters_dict = dict(model.named_parameters())
 
         param_to_name_index = {param: name for name, param in named_parameters_dict.items()}

{x_evolution-0.1.19.dist-info → x_evolution-0.1.23.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-evolution
-Version: 0.1.19
+Version: 0.1.23
 Summary: x-evolution
 Project-URL: Homepage, https://pypi.org/project/x-evolution/
 Project-URL: Repository, https://github.com/lucidrains/x-evolution

x_evolution-0.1.23.dist-info/RECORD

@@ -0,0 +1,6 @@
+x_evolution/__init__.py,sha256=XcwXJgIMPnCWGfGws3-vKgoR_7IfVslJBtiMvmEeSg0,57
+x_evolution/x_evolution.py,sha256=lvN3ePqD6a5dW1gOv0d1I9yQ4rdv6OuIVvKvXa0yRBM,19126
+x_evolution-0.1.23.dist-info/METADATA,sha256=pYkzZj10t68WjpyFOSMzd8lhRYL9FNGRsHjlcLyP9zw,5728
+x_evolution-0.1.23.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+x_evolution-0.1.23.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
+x_evolution-0.1.23.dist-info/RECORD,,

x_evolution-0.1.19.dist-info/RECORD

@@ -1,6 +0,0 @@
-x_evolution/__init__.py,sha256=XcwXJgIMPnCWGfGws3-vKgoR_7IfVslJBtiMvmEeSg0,57
-x_evolution/x_evolution.py,sha256=C70fX5PdLx1VLBXe1DYGBIOExucvLFbKDWiofBMrvQA,18860
-x_evolution-0.1.19.dist-info/METADATA,sha256=7xr4uJpkmkkg0bi9eByDktqRvm_IIcvHJcLAJUjB2QM,5728
-x_evolution-0.1.19.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-x_evolution-0.1.19.dist-info/licenses/LICENSE,sha256=1yCiA9b5nhslTavxPjsQAO-wpOnwJR9-l8LTVi7GJuk,1066
-x_evolution-0.1.19.dist-info/RECORD,,