How to train custom made dataset for beam engineering problem?

Hi!

Greetings! In my machine learning problem, I have an input, which is a Gaussian beam, The input Gaussian beam has to pass through different layers and finally produce a beam of desired shape. My model looks something like this,

E= A* torch.exp(-((x)**2+(y)2)/(wa2)); # This is a analytical formula for gaussian beam
G=torch.fft.fft2(E) # Taking Fourier transform of it.

Propagate through the 1st layer

elayer1=torch.fft.ifft2(torch.fft.ifftshift((torch.fft.fftshift(G))*H0)); # Multiplying it with the transfer function

Multiplying Elayer1 with phase1 values

elayer1=elayer1torch.exp(1j(Phase1));
predictedbeam=torch.abs(elayer1)**2

In the first line, I have given the Gaussian beam. I want to train this model, with many number of gaussian beams that are tilted slightly (E= A* torch. exp(-((x-a)*2+(y-b)2)/(wa2)). So that the model becomes robust even if there is a slight tilt in the Gaussian beam position. I have made a custom-made dataset. I classified the Gaussian beam based on how they are titled (E= A torch. exp(-((x-a)**2+(y-b)2)/(wa2))). Now I want to train them. In the place of E, in the code, I want to give a folder that has many Gaussian beams. I don’t know how to go about it. Can anyone kindly help me with this?