Fix missing import.
Merge pull request #22 from Debilski/fix-import

Fix missing import.