torch.jit.script fails on (time unrolled) norse modules

[cpehle]

Describe the bug

This is not strictly speaking a bug, but at the moment torch.jit.script fails on the following
network with

RuntimeError: 
Module 'LIF' has no attribute 'state_fallback' (This attribute exists on the Python module, but we failed to convert Python type: 'method' to a TorchScript type. Only tensors and (possibly nested) tuples of tensors, lists, or dictsare supported as inputs or outputs of traced functions, but instead got value of type method.. Its type was inferred; try adding a type annotation for the attribute.):
  File "/home/cpehle/work/norse/norse/torch/module/snn.py", line 210
    def forward(self, input_tensor: torch.Tensor, state: Optional[Any] = None):
        state = state if state is not None else self.state_fallback(input_tensor)
                                                ~~~~~~~~~~~~~~~~~~~ <--- HERE

I think it is worthwhile changing the API so that this goes through. I will potentially need this in the near future to support hardware backends via this mechanism.

To Reproduce

import torch
from norse.torch import LIF, LI

model = torch.nn.Sequential(
    torch.nn.Conv2d(1, 20, 5, 1),
    LIF(),
    torch.nn.Conv2d(20, 50, 5, 1),
    LIF(),
    torch.nn.Flatten(),
    torch.nn.Linear(800, 10),
    LI(),
)

torch.jit.script(model)

Expected behavior

The module can be jitted.

[Jegp]

Did you try to provide the state? This should only happen if the state is undefined.

[cpehle]

We should have tests for this and document how to jit.