I'm working on a project where I need my model to predict a sequence of `n` 3x3 matrices given an input sequence of `n` 3x3 matrices for a physical simulation. (the `n`'s are all perfect squares e.g. I only need to consider sequences of 9, 16, 25, 36, etc. matrices all the way to 324). I need the model to approximate the mathematical relationship between the input and output (which can be computationally expensive for large `n`, hence the need for an ML model in the first place).

However, the problem I have is that my dataset only contains input/output values for sequences of 9, 16, ... 81 matrices (after that it becomes hard to compute).

I am currently a novice in machine learning and just know the basics of Pytorch, but so far I have trained a simple neural network with alternating Linear and ReLU layers for the `n=9` case, with error ~ `10^-4`. However, now I want to expand to the general case, but I am not sure of the best way to do so.

What I am thinking is to make my neural network big enough for the largest case `n=324` and use some kind of RNN to make the predictions accurate for large `n` even though my dataset does not cover all those values.

I just wanted to ask if this is the right approach or if there is a better way to transfer the information learned from small dataset sizes to make predictions for larger cases, and also what type of RNN I should use.

Thank you!