Fix memory allocation