Add missing import Fix #36

Bug introduced in #25