Improve Training Loop for Autoencoder and Classifier Models

[Ankitaghavate]

The current train.py handles training for autoencoder, classifier partial, and classifier full models with masking, tensor reshaping, and validation after each epoch. The code is repetitive, uses hardcoded values, and relies on print statements for logging. It needs refactoring for modularity, consistent type hints, dynamic logging, better error handling, and cleaner saving of best models, making the training process more maintainable and robust.

[Abhi1366-create]

Hi, thanks for opening this issue. I went through train.py and see that the training logic for autoencoder, classifier partial, and classifier full shares a lot of repeated structure (epoch loop, batch loop, optimizer steps, logging, and validation).

As a first incremental step, would it make sense to replace the print-based logging with Python’s logging module and centralize the logging configuration? That could be done without touching the training logic itself. Happy to start with a small PR if that aligns with the plan.