heavyball 0.24.0__tar.gz → 0.24.1__tar.gz

{heavyball-0.24.0 → heavyball-0.24.1}/PKG-INFO

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

{heavyball-0.24.0 → heavyball-0.24.1}/heavyball/utils.py

@@ -134,7 +134,7 @@ def dim_merger(grad, max_precond_dim, split: bool = False):
     if curr_shape > 1 or len(new_shape) == 0:
         new_shape.append(curr_shape)
 
-    new_grad = grad.view(new_shape)
+    new_grad = grad.reshape(new_shape)  # needs to be .reshape() due to channels_last
     if not split:
         return new_grad
 

{heavyball-0.24.0 → heavyball-0.24.1}/heavyball.egg-info/PKG-INFO

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

{heavyball-0.24.0 → heavyball-0.24.1}/setup.py

@@ -10,7 +10,7 @@ setuptools.setup(
     name='heavyball',
     license='BSD',
     description='Efficient optimizers',
-    version='0.24.0',
+    version='0.24.1',
     long_description=README,
     url='https://github.com/clashluke/heavyball',
     packages=setuptools.find_packages(),

{heavyball-0.24.0 → heavyball-0.24.1}/test/test_channels_last.py

@@ -11,6 +11,7 @@ from heavyball.utils import clean, set_torch
 from torch import nn
 from torch._dynamo import config
 
+heavyball.utils.compile_mode = 'default'
 config.cache_size_limit = 128
 
 
@@ -29,14 +30,14 @@ def test_foreach(opt, size, depth: int, iterations: int = 32, outer_iterations:
         losses.append([])
 
         for i in range(outer_iterations):
-            model = nn.Sequential(*[nn.Conv2d(size, size, 1) for _ in range(depth)]).cuda()
+            model = nn.Sequential(*[nn.Conv2d(size, size, 3) for _ in range(depth)]).cuda()
             if is_channels_last:
                 model.to(memory_format=torch.channels_last)
 
             o = get_optim(opt, model.parameters(), lr=1e-3, weight_decay=1e-4, warmup_steps=16)
 
             for _ in range(iterations):
-                loss = model(torch.randn((1024, size, 1, 1), device='cuda')).square().mean()
+                loss = model(torch.randn((1024, size, 4, 4), device='cuda')).square().mean()
                 loss.backward()
                 o.step()
                 o.zero_grad()