Suggested to add a check for the return value of Malloc

[dududuguo]

test_gpt2.c:55-58 train_gpt2.c:484 \531\ 634 Many implementations of malloc should include a checker to verify successful memory allocation. Suggested to add NULL checks after each malloc call and handle the memory allocation failures gracefully. Thank you!