How to train a network, using the gradient of the loss of the function of the output of the neural network?

I have some data with complex values. What I do (want to do) is essentially this:

Say, I have a complex number z = a + ib. I apply the neural network on 'a' and then combine it with 'b', apply another function on the new complex number and then compare the imaginary part of this new number with another number f, which is obtained from another dataset.

Explanation of what I'm doing

Here's how I try to implement it on python using Pytorch:

class NN(nn.Module):
    def __init__(self,size):
        
        super(NN, self).__init__()                  
        self.linear1 = nn.Linear(size, 2*size)
        self.Softplus = nn.Softplus()
        self.linear2 = nn.Linear(2*size, size)
    

    
    def forward(self,x):
        
        x = self.linear1(x)
        x = self.Softplus(x)
        x = self.linear2(x)

        
        return x   

Also, I am using the Adam algorithm for the optimizer and the L1Loss function for the loss. Here's what my training function looks like

 def train(dataloader, model, loss_fn, optimizer):   
    
    size = len(dataloader.dataset)
    model.train()
    num_batches = len(dataloader)
    train_loss = 0
    
    for batch, (X, y,z) in enumerate(dataloader):
        X, y, z = X.to(device), y.to(device), z.to(device) 
        
        ftr = model(X)  #Applying the neural network
        pred = func(ftr,y)   #Applying the function on a and b
        
        loss = loss_fn(pred,z)  #computing loss (d,f)
        loss = Variable(loss, requires_grad = True)
        
        train_loss += loss.item()

        
        # Backpropagation
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
    
    train_loss /= num_batches
    return train_loss

The issue is that whenever I do this, my loss seems to be constant and there's a huge gap between the validation and training loss, with both varying very little to 0 for all epochs.

I believe this is due to some issues with computing gradients and the losses not being directly associated with the weights. But I am not sure - any ideas?