I think it would be useful to have a BaseModule machine wrapper around nn.Module. Similarly to : https://github.com/YannDubs/machine/blob/extrapolator/seq2seq/util/base.py#L14 . Just removing the functions for visualisation as we will be using callbacks to separarte visualisation from machine.