Added some code to run.py