Can you create a library based file out of it, such that I can just import and call?

[SaumyaBhandari]

Try making it in an object oriented fashion such that, I can import it... pass some parameters on how much epoch I want to train, batch size and it should just run