Minor fix for a wrong import.