Currently, there is a some code duplication in AptaTrans lightning trainig code: both classes in aptatrans._model_lightning. Specifically, the logging calls in training_step and test_step are redundant and can be moved to the _step(...) method.
This issue is better addressed when #190 is merged, since tests and some the interface have been slighly changed. The code duplication issue is still relevant though. Therefore parking this for now.