Fix bugs after initial loading test