Unable to calculate differential using torch.autograd.grad

Hello all, I’m trying to calculate the differential of the output with respect to one of the column in the dataset. I am not able to produce the differential. It is a regression problem and the network has two outputs. There are more than 2000 columns and only 1 column is required in the calculation of differential.

Error produced is as follows:

RuntimeError: One of the differentiated Tensors appears to not have been used in the graph. Set allow_unused=True if this is the desired behavior.

Model :

class Model2(nn.Module):

def __init__(self, input_dim):

        super(Model2, self).__init__()
        self.layer1 = nn.Linear(input_dim, 512)
        self.layer2 = nn.Linear(512, 256)
        self.layer3 = nn.Linear(256, 128)
        self.layer4 = nn.Linear(128, 64)
        self.layer5 = nn.Linear(64, 2)
        
    def forward(self, x):
        x = F.relu(self.layer1(x))
        x = F.relu(self.layer2(x))
        x = F.relu(self.layer3(x))
        x = F.relu(self.layer4(x))
        x = self.layer5(x)
        return x

Training:

EPOCHS = 1000
fourth_loss_train = np.zeros((EPOCHS,))
fourth_loss_val = np.zeros((EPOCHS,))
fourth_train_mape = np.zeros((EPOCHS,))
fourth_val_mape = np.zeros((EPOCHS,))

x1 = X_train[:,-2].requires_grad_(True)
x2 = 1 - x1 #.requires_grad_(True)

for epoch in tqdm.trange(EPOCHS, position=0, leave=True):

model_PINN_5.train()

y_pred = model_PINN_5(X_train)
loss1 = loss_fn(y_pred, y_train)

d_lngamma1_x1 = torch.autograd.grad(y_pred[:,0], x1, torch.ones_like(y_pred[:,0]), create_graph=True)[0]
d_lngamma2_x2 = torch.autograd.grad(y_pred[:,1], x2, torch.ones_like(y_pred[:,1]), create_graph=True)[0]

physics = x1@d_lngamma1_x1 # - x2@d_lngamma2_x2
loss2 = (1e-4)*torch.mean(physics**2)

loss = loss1 + loss2

fourth_loss_train[epoch] = loss.item()

# Zero gradients
loss.backward()
optimizer.step()
optimizer.zero_grad()


# Calculate train accuracy
fourth_train_mape[epoch] = mean_absolute_percentage_error(y_train, y_pred)

with torch.no_grad():
    model_PINN_5.eval()
    
    y_pred = model_PINN_5(X_val)
    val_loss = loss_fn(y_pred, y_val)
    fourth_loss_val[epoch] = val_loss
    
    fourth_val_mape[epoch] = mean_absolute_percentage_error(y_val, y_pred)