Merge pull request #369 from TBBristol/heavyball

Thank you

Author: jsuarez5341

setup.py

@@ -258,7 +258,7 @@ def run(self):
         'rich_argparse',
         'imageio',
         'pyro-ppl',
-        'heavyball',
+        'heavyball<2.0.0',
         'neptune',
         'wandb',
     ]