Fix to null reference issue