I don't know which numerical optimization OP is trying to solve, but one major weakness of this paper is that their method requires two solver calls per instance per epoch... Training time might quickly become intractable.

OP should have a look at other methods that aim to solve their problem efficiently, such as https://arxiv.org/abs/2112.03609 or recently https://arxiv.org/abs/2203.16067

You have this survey about ML for Combinatorial Optimization from Bengio, Lodi, and Provost.

OpenAI's paper about a robot hand learning to solve a rubik's cube

Also check-out anything mentioning neurosymbolic AI, which aims to combine neural network learning with logic-based reasoning.

Gary Marcus wrote an extensive note on the subject that I recommend as well