That will definitely fix the grader error. But did that fix also cure your problem with the kernel dying when you run the training? If not, one thing to check is whether you’ve added any print statements that cause the output of the training to be more voluminous that it would be by default. That can cause issues if the memory image of the notebook gets too large.