heavyball 0.18.6__py3-none-any.whl → 0.18.7__py3-none-any.whl

heavyball/utils.py

@@ -843,12 +843,13 @@ class PSGDBase(StatefulOptimizer):
             psgd_update_precond(Q, self.state_(p)["exprs"], torch.randn_like(grad), grad, precond_lr, self._tiny)
             update_fn(oq, Q)
 
-        for g, q in zip(grad_list, original_q if original_q else q_list):
-            if g.dim() > 1:
-                if store_triu_as_line:
-                    psgd_balance_Q([q_ for _, q_ in q])
-                else:
-                    psgd_balance_Q(q)
+        if self.should_update(group, self.balance_probability, "balance_prob"):
+            for g, q in zip(grad_list, original_q if original_q else q_list):
+                if g.dim() > 1:
+                    if store_triu_as_line:
+                        psgd_balance_Q([q_ for _, q_ in q])
+                    else:
+                        psgd_balance_Q(q)
 
 
 def precond_update_prob_schedule(max_prob=1.0, min_prob=0.03, decay=0.001, flat_start=250):

{heavyball-0.18.6.dist-info → heavyball-0.18.7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: heavyball
-Version: 0.18.6
+Version: 0.18.7
 Summary: Efficient optimizers
 Home-page: https://github.com/clashluke/heavyball
 Author: Lucas Nestler

{heavyball-0.18.6.dist-info → heavyball-0.18.7.dist-info}/RECORD

@@ -16,9 +16,9 @@ heavyball/precond_schedule_sfpsoap.py,sha256=vq7jd302refKPa_9X2lkOTOtCCcTBVByPdo
 heavyball/psgd_kron.py,sha256=u46dorOUXx-do1IYeno2wj-6l1zYKMQQC-N2Zr2PzLI,5476
 heavyball/pure_psgd.py,sha256=iUy7mMKWxwNiVUMYrQ7SBnreu3t_XSbnhTW3a1yw4m0,4835
 heavyball/schedule_free_palm_foreach_soap.py,sha256=zkcikH5wWbzq4kOrmBjilvY3iWzuUddcv2HNEPKr3MI,6366
-heavyball/utils.py,sha256=UN0nqXSF-NkRtvCZUdlrs4sRJzZy1v89BkixasQ-bZg,31304
-heavyball-0.18.6.dist-info/LICENSE,sha256=CGdGJim64YifGmUVPaeyRsxkvyExtClswhRNIp8FY_U,1322
-heavyball-0.18.6.dist-info/METADATA,sha256=Grc4Ok9z1BdjtZ7XqRHedKLGNNMcyAU8-QpcGZNl8Hk,11810
-heavyball-0.18.6.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-heavyball-0.18.6.dist-info/top_level.txt,sha256=SzCxSVg_qCUPA4kZObW3Zyo4v-d_mMOD-p7a-WXTl2E,10
-heavyball-0.18.6.dist-info/RECORD,,
+heavyball/utils.py,sha256=lKIV11qvlHITK7lwaScGbP1ryCmInse9Fe64t0OBmQQ,31408
+heavyball-0.18.7.dist-info/LICENSE,sha256=CGdGJim64YifGmUVPaeyRsxkvyExtClswhRNIp8FY_U,1322
+heavyball-0.18.7.dist-info/METADATA,sha256=KUYpVlwytyMmQBuby0Jf1WaklYdc2GPddiMAqyGKzsM,11810
+heavyball-0.18.7.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+heavyball-0.18.7.dist-info/top_level.txt,sha256=SzCxSVg_qCUPA4kZObW3Zyo4v-d_mMOD-p7a-WXTl2E,10
+heavyball-0.18.7.dist-info/RECORD,,