Fixed type and removed unused imports
Closes #8