fixed the broken code