Add missing import

This is used in `RawPositionsFeaturizer`