Updated main.py to avoid Travis hanging