Hi. I am quite new to Pytorch and learning it by trying out some example notebooks. The one I am busy with now involves an unsupervised neural network for solving an eigenvalue problem in Quantum Mechanics, a one-dimensional Schrodinger equation with an infinite square-well potential. The  ipynb notebook is provided here:
eigeNN/BothBounds_Infinite.ipynb at master · henry1jin/eigeNN · GitHub.
In the context of this code, the neural network represents the eigenfunction or solution of the differential equation. The two inputs include one adjustable scalar, E, representing the eigenvalue (which is learned during training) and one independent variable, t, of the differential equation. The neural network’s one output is the eigenfunction, f(t). In code, this is:
ORIGINAL:
# Define the sin() activation function
class mySin(torch.nn.Module):
    @staticmethod
    def forward(input):
        return torch.sin(input)

class qNN1(torch.nn.Module):
    def __init__(self, D_hid=10):
        super(qNN1,self).__init__()

        # Define the activation 
        self.actF = mySin()
        
        # define layers
        self.Ein    = torch.nn.Linear(1,1)
        self.Lin_1  = torch.nn.Linear(2, D_hid)
        self.Lin_2  = torch.nn.Linear(D_hid, D_hid)
        self.out    = torch.nn.Linear(D_hid, 1)

    def forward(self,t):
        In1 = self.Ein(torch.ones_like(t))
        L1 = self.Lin_1(torch.cat((t,In1),1))
        h1 = self.actF(L1)
        L2 = self.Lin_2(h1)
        h2 = self.actF(L2)
        out = self.out(h2)
        return out, In1

The complete example code is able to compute the eigenvalues and eigenfunctions for the lowest bound states using a scanning mechanism. I’m trying out a slight addition to the code which, in theory, shouldn’t change the result. That is, I’m considering the eigenfunctions and eigenfunction as complex-valued scalars and not real scalars as given in the original code. This means I needed to include two inputs and two outputs for the real and imaginary parts of E and f(t), respectively. Here, I am not sure if I implemented that correctly:
MODIFIED:
# Define the sin() activation function
class mySin(torch.nn.Module):
    @staticmethod
    def forward(input):
        return torch.sin(input)

class qNN1(torch.nn.Module):
    def __init__(self, D_hid=10):
        super(qNN1,self).__init__()

        # Define the activation
        self.actF = mySin()
        
        # define layers
        self.E1in    = torch.nn.Linear(1,1)
        self.E2in    = torch.nn.Linear(1,1)
        self.Lin_1  = torch.nn.Linear(3, D_hid) # three inputs
        self.Lin_2  = torch.nn.Linear(D_hid, D_hid)
        self.out    = torch.nn.Linear(D_hid, 2) # two outputs

    def forward(self,t):
        In1 = self.E1in(torch.ones_like(t))
        In2 = self.E2in(torch.ones_like(t))
        L1 = self.Lin_1(torch.cat((t,In1,In2),1)) # In2 added
        h1 = self.actF(L1)
        L2 = self.Lin_2(h1)
        h2 = self.actF(L2)
        out = self.out(h2)
        return out, In1, In2

Boundary conditions are imposed as physical constraints on the neural network model and the differential equation itself is embedded in the loss function via autodifferentiation.
ORIGINAL:
def parametricSolutions(t, nn, t0, x1):
    # parametric solutions 
    N1,N2 = nn(t)
    dt =t-t0
    f = (1-torch.exp(-dt))*(1-torch.exp(dt-1)) # this parametric function ensure the neural network model satisfies the boundary conditions by construction
    psi_hat  = x1  + f*N1
    return psi_hat

def dfx(x,f):
    # Calculate the derivative with auto-differention
    return grad([f], [x], grad_outputs=torch.ones(x.shape, dtype=dtype), create_graph=True)[0]

def hamEqs_Loss(t,psi, E):
    # one-dimensional Schrodinger equation with an infinite square-well potential
    psi_dx = dfx(t,psi)
    psi_ddx= dfx(t,psi_dx)
    f = psi_ddx/2 + E*psi
    L  = (f.pow(2)).mean(); #part of the total loss function due to differential equation "residual"
    return L

Continuing with the assumption I made above concerning complex scalars, the changes I made to cater for that are:
MODIFIED:
def parametricSolutions(t, nn, t0, x1):
    # parametric solutions 
    N1 = nn(t)[0] # first neural network output
    N2 = nn(t)[1] # second neural network output
    dt =t-t0
    f = (1-torch.exp(-dt))*(1-torch.exp(dt-1))
    psi1_hat = x1  + f*N1
    psi2_hat  = x1 + f*N2
    return psi1_hat, psi2_hat

def dfx(x,f):
    # Calculate the derivative with auto-differention
    return grad([f], [x], grad_outputs=torch.ones(f.shape, dtype=dtype), create_graph=True)[0]

def hamEqs_Loss(t,psi1, psi2, E1, E2):
    # one-dimensional S.E.
    psi1_dx = dfx(t,psi1)
    psi1_ddx= dfx(t,psi1_dx)
    psi2_dx = dfx(t,psi2)
    psi2_ddx= dfx(t,psi2_dx)
    f  = (psi1_ddx + psi2_ddx*1j)/2 + (E1 + E2*1j)*psi
    f1 = torch.real(f)
    f2 = torch.imag(f)
    L  = (f1.pow(2)).mean() + (f2.pow(2)).mean(); 
    return L

My two main concerns here are with the syntax for indexing the two outputs of the neural network and the change I made to grad_outputs in def dfx(x,f) . I’m not sure whether I’ve correctly assigned the N1 and N2 of def parametricSolutions(t, nn, t0, x1): to the two neural network outputs separately. Another change I am worried about is the one I made to the shape of grad_outputs from x.shape to f.shape to make dfx compatible with a two output neural network.
All this comes together in the training stage of the code:
ORIGINAL:
# Train the NN
def run_Scan_finitewell(t0, tf, x1, neurons, epochs, n_train,lr, minibatch_number = 1):
    fc0 = qNN1(neurons)
    ...
    SE_loss_history = []
    ...
    for tt in range(epochs): 
    ...
#  Network solutions 
            nn, En = fc0(t_mb) #t_mb related to the training inputs   
            psi  = parametricSolutions(t_mb, fc0, t0, x1) 
            Ltot = hamEqs_Loss(t_mb, psi, En)
            SE_loss_history.append(Ltot) 

My version with real and imaginary parts:
MODIFIED:
# Train the NN
def run_Scan_finitewell(t0, tf, x1, neurons, epochs, n_train,lr, minibatch_number = 1):
    fc0 = qNN1(neurons)
    ...
    SE_loss_history = []
    ...
    for tt in range(epochs): 
    ...
#  Network solutions 
            nn, E1n, E2n = fc0(t_mb) #t_mb related to the training inputs   
            psi1, psi2  = parametricSolutions(t_mb, fc0, t0, x1) 
            Ltot = hamEqs_Loss(t_mb, psi1, psi2, E1n, E2n)
            SE_loss_history.append(Ltot) 

After making these changes I hoped to generate the same results as the original code except that the imaginary parts would be zero. However when I run my modified version with the same number of epochs (120K) as in the original example, the RAM is filled up faster than in the original code and the session crashes before the training has ended. I’d really appreciate suggestions from anyone regarding this.
Additionally, I would like to check if I’ve correctly added two neural networks outputs and properly indexed them in the other parts of the code. My modified version runs without errors and the only apparent issue is the aforementioned crashing after several epochs. Please assist.