Added some comments to main.py

Signed-off-by: Archit Shukla <architshukla2@gmail.com>