Repository navigation
Improve Training Loop for Autoencoder and Classifier Models #17
Description
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.
Activity
Abhi1366-create commentedon Feb 8, 2026
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.
feat: Improve training loop for autoencoder and classifier models (fixes
VinayTodkar commentedon Aug 1, 2026
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.
higgsboson1710 commentedon Sep 9, 2026
Can u please assign me . I can look after it !!