Fixed bug caused by missing variables after refactor