Fix import mistake