change to main.py

added some code