Small Bug Fix

Fixed a small bug and added a function to main.py