Fix memory allocation issue