Update main.py

Make main.py in step 4 consistent with main.py in the other steps
Merge pull request #1132 from rogererens/patch-1

Update main.py