Pytorch error when launching two distinct backward

I am building a simple autoencoder followed by an MLP neural nets. Regarging the autoencoder I am not running into any problem

# ---- Prepare training set ----
x_data = train_set_categorical.drop(["churn"], axis=1).to_numpy()
labels = train_set_categorical.loc[:, "churn"].to_numpy()

dataset = TensorDataset(torch.Tensor(x_data), torch.Tensor(labels) )

loader = DataLoader(dataset, batch_size=127)

# ---- Model Initialization ----
model = AE()
  
# Validation using MSE Loss function
loss_function = nn.MSELoss()
  
# Using an Adam Optimizer with lr = 0.1
optimizer = torch.optim.Adam(model.parameters(),
                             lr = 1e-1,
                             weight_decay = 1e-8)

epochs = 50
outputs = []
losses = []
for epoch in range(epochs):
    for (image, _) in loader:

        # Output of Autoencoder
        embbeding, reconstructed = model(image)

        # Calculating the loss function
        loss = loss_function(reconstructed, image)

        # The gradients are set to zero,
        # the the gradient is computed and stored.
        # .step() performs parameter update
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # Storing the losses in a list for plotting
        losses.append(loss)
        if epoch == 49:
            outputs.append(embbeding)

But then I am feeding an MLP with the outcome of the autoencoder and this is where things starts to fail

class Feedforward(torch.nn.Module):
    def __init__(self):
        super().__init__()
        
        self.neural = torch.nn.Sequential(
            torch.nn.Linear(33, 260),
            torch.nn.ReLU(),
            torch.nn.Linear(260, 450),
            torch.nn.ReLU(),
            torch.nn.Linear(450, 260),
            torch.nn.ReLU(),
            torch.nn.Linear(260, 1),
            torch.nn.Sigmoid()
        )
        
    def forward(self, x):
        outcome = self.neural(x.float())
        return outcome
    
modelz = Feedforward()
criterion = torch.nn.BCELoss()
opt = torch.optim.Adam(modelz.parameters(), lr = 0.01)

modelz.train()
epoch = 20
for epoch in range(epoch):
    opt.zero_grad()
    # Forward pass
    y_pred = modelz(x_train)
    # Compute Loss
    loss_2 = criterion(y_pred.squeeze(), torch.tensor(y_train).to(torch.float32))
   
    #print('Epoch {}: train loss: {}'.format(epoch, loss.item()))
    # Backward pass
    loss_2.backward()
    opt.step()

I get the following error:

RuntimeError: Trying to backward through the graph a second time, but the saved intermediate results have already been freed. Specify retain_graph=True when calling .backward() or autograd.grad() the first time.

Of course I have tried to add "retain_graph=True" to both backwards or only the first one but it does not seem to solve the problem. If I launch both code independently from another It works but as a sequence I don't know why but it is not.

Solution 1:

You should be able to disconnect the output of the auto-encoder from the model by calling embbeding.detach(), before appending it to outputs.