Grading error (code is not optimised)

When I submit the code, it takes an hour to grade and returns an unexpected error saying that i need to optimize my code. I’ve stuck on this for hours. Please help

Code not optimized usually points to very big model architecture / training for a lot of epochs. This does not mean that your solution is incorrect. It means that the a smaller NN would suffice to solve the problem and achieve the desired thresholds.

My code uses 10 epochs to reach 99.8% accuracy. Should I decrease the number of epochs still?

Have you tried making your model smaller and submitting again?