Update trainer tests.
Add tests for dataset class.
Slice 3d datasets along axis 1 by default.
Refactor dataset tests.