Added import in code. Was causing issues

Hi I might have submitted the pull request using a wrong branch so there was an import statement missing. Now added import.