fixed a bug in main.py