Fix import

Late night commit, duh.