C++ frontend examples on training with customized dataset

Dear Pytorch users,
I am coming up with a pytorch enabled surrogate reference for my so-called multi-scale simulation framework, that inherits the material properties from it’s micro-structure considerations. The idea here is to use my macro-model in conjunction with the surrogate model that effectively substitutes the finer scale material constitutive behavior. With a careful selection of the training dataset in hand, i am looking at options in pytorch based c++ frontend to train the surrogates with this dataset and reference them with the target variables the macro model is throwing out.

The training dataset has the following simple structure that comprises the deformation gradient and stresses. You might find several of these lines for many different simulations that constitute this dataset.
F1 F2 F3 F4 F5 F6 F7 F8 F9 S1 S2 S3 S4 S5 S6 S7 S8 S9

The macro-model will examine the trained set for the macro-model deformation gradients of this type and returns the stresses
F1 F2 F3 F4 F5 F6 F7 F8 F9

We made some progress on this imposing the dataset on a NN model creating a torch script module following this approach Torch Script Module. Later this turned out to be an incorrect selection, as the length of the training dataset and evaluation set is required to be identical, without which the approach will be infeasible.

The pytorch examples here has references to some good examples on usage of NN models within an optimization setting, but has limited focus on creating the kind of datasets to be used within this framework. The simplest way one can compile this dataset is to put all these 18 numbers in a sequence of vectors in a CSV file and convert them to the data structure the data loader demands. The better and the recommended way is to use a MNIST dataset format which has these numbers encoded in a hand written format that has a natural connection with these examples. Given these I am not very clear on the following 1) How good is the selection of MNIST dataset, that has a better representation of the training dataset mentioned above 2) If all seems well, how this dataset can be pushed in to a MNIST dataset (any references to examples are appreciated).

Thanks a lot for looking at this case. Any suggestions/leads on handling this issue will be much appreciated.
Thanks and regards,
Arun