AttributeError: 'BatchNorm2d' object has no attribute 'output'

I post the full error below. The MWE is a bit long (currently hundreds of lines) and I am still working on it, but is there any specific direction I should be looking at given this error? It looks like Batchnorm is somehow mixed up in the gradient calculation (judging from the error message)?

```
Traceback (most recent call last):
File "/Users/qiyaowei/DEQ-BNN/mwe.py", line 575, in
model(torch.rand(1,3,32,32)).sum().backward()
File "/Users/qiyaowei/miniconda3/envs/jax/lib/python3.8/site-packages/torch/_tensor.py", line 363, in backward
torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)
File "/Users/qiyaowei/miniconda3/envs/jax/lib/python3.8/site-packages/torch/autograd/__init__.py", line 173, in backward
Variable._execution_engine.run_backward( # Calls into the C++ engine to run the backward pass
File "/Users/qiyaowei/miniconda3/envs/jax/lib/python3.8/site-packages/torch/utils/hooks.py", line 110, in hook
res = user_hook(self.module, grad_input, self.grad_outputs)
File "/Users/qiyaowei/miniconda3/envs/jax/lib/python3.8/site-packages/backpack/__init__.py", line 209, in hook_run_extensions
backpack_extension(module, g_inp, g_out)
File "/Users/qiyaowei/miniconda3/envs/jax/lib/python3.8/site-packages/backpack/extensions/backprop_extension.py", line 127, in __call__
module_extension(self, module, g_inp, g_out)
File "/Users/qiyaowei/miniconda3/envs/jax/lib/python3.8/site-packages/backpack/extensions/module_extension.py", line 97, in __call__
delete_old_quantities = not self.__should_retain_backproped_quantities(module)
File "/Users/qiyaowei/miniconda3/envs/jax/lib/python3.8/site-packages/backpack/extensions/module_extension.py", line 162, in __should_retain_backproped_quantities
is_a_leaf = module.output.grad_fn is None
File "/Users/qiyaowei/miniconda3/envs/jax/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1185, in __getattr__
raise AttributeError("'{}' object has no attribute '{}'".format(
AttributeError: 'BatchNorm2d' object has no attribute 'output'
```