Update main.c

Added some memory handling