Add tests for inference / training

[cpehle]

By checking a small model trained on a simple task, it would become easier to test training and inference. It would also help test TorchScript support.

[Jegp]

Is this closed with #123 @cpehle ?

[cpehle]

Is this closed with #123 @cpehle ?

The issue is a bit unspecific, but what I meant was a test which does end-to-end training on a toy dataset which doesn't blow up the CI budget.

[Jegp]

How about using the odd-firing model in https://github.com/norse/notebooks/blob/master/parameter-learning.ipynb?
It's sufficiently complex to have to be learned and it shouldn't steal too many resources.