Memory allocation bug fix