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.

[VinayTodkar]

Hi, I would like to work on this issue. I plan to refactor the training loops for the autoencoder and classifier models, reduce code duplication,

replace hardcoded values with configurable parameters, introduce structured logging, improve type hints and error handling, and clean up

best-model checkpointing while ensuring the current training behavior is preserved. Please let me know if this approach aligns with the

project's expectations.