Adding a training loop

[SantoshParit]

pushing the training loop and confusion matrix code to main.

@Sohambasu07 @sushil579 @kbpranav @deepukr007