Adding more flags to main.py

[JonnaMat]

Hi,

I would like to contribute and make main.py a proper training script.

Adding, for example,

if __name__ == "__main__":
    main()

and restructuring the code accordingly.

and some flags (for both, HP optimization (--bacth-size/-b, --epochs, --momentum, ...) and general ease of use (--device, --data-path, ...) - generally, sticking to "naming conventions" in the torchvision reference training scripts.

Let me know how contribution works in this project.

Thanks!