There are a few libraries/frameworks that one can use and allow to reuse the same code for datasets, logging, training loop etc.... . E.g. Lightning or Scenic. Maybe you can use one of these or at least get some inspiration for your own code.