Add forgotten import statement

When I refactored, this got lost.