Fix null reference bug