Bug fix (Still not the best solution)