Allow any model architecture

[ClashLuke]

No description provided.

[ClashLuke]

Works well. example notebook

[dvruette]

Looks great! Ideally we would update at least the test/autoencoder.py example (and the README) for future reference.

I've left some comments where I didn't quite understand.

Another thing I've noticed yesterday is that torch.autograd.grad doesn't work with integer tensors as inputs, which happens e.g. when using nn.Embedding. Not sure if this is easily resolvable, so we don't need to address it now. But something to keep in mind for the future.

[dvruette]

While training seems to work with Lightning, the standard training loop breaks for some reason: https://colab.research.google.com/drive/1hNLavl5jYgf7-DxfmnTucyCAcMHS4y4c?usp=sharing

Don't really know what's going on, the only difference is the training loop AFAICT.

[ClashLuke]

We get roughly the same convergence with an outer learning rate of 0.01. Look at this copy of your notebook.
Does the lightning trainer do some gradient scaling internally?