Fixed bugs in main.py
Better model params