Add missing imports
Merge pull request #14 from fmr/fix/add_missing_import

Add missing imports