Getting error from running README example...

[slerman12]

autograd_hacks.compute_grad1() TypeError: compute_grad1() missing 1 required positional argument: 'model'

[slerman12]

I also get this error when I add model as an argument in that same README example:

AttributeError: 'Parameter' object has no attribute 'grad1'

[slerman12]

Oh I think it's because my model consists of submodules:

Model(
  (torso): MLP(
    (model): Sequential(
      (0): Linear(in_features=5, out_features=64, bias=True)
      (1): Tanh()
      (2): Linear(in_features=64, out_features=64, bias=True)
      (3): Tanh()
    )
  )
  (head): DetachedScaleGaussian(
    (loc_layer): Sequential(
      (0): Linear(in_features=64, out_features=1, bias=True)
      (1): Tanh()
    )
  )
)

Any idea what I can do?