Improve metrics logging

[pranshu-raj-211]

Fixes to be done:

• Add logging for test metrics (a sort of report of classification metrics)
• Add train-test loss curves to qualitatively measure performance and fit
• Create and analyze confusion matrix
• Log all plots

[pranshu-raj-211]

Use this to improve logging wandb docs. Modularize metric checking.