Fix missing imports

Oopsie! :rose: