Fixing import bugs
Merge pull request #2 from lemurell/master

Fixed imports in main.