Hi, I’ve written a custom ML model (a quantum circuit), however when it comes to training I am getting the following error:

RuntimeError: Trying to backward through the graph a second time (or directly access saved variables after they have already been freed). Saved intermediate values of the graph are freed when you call .backward() or autograd.grad(). Specify retain_graph=True if you need to backward through the graph a second time or if you need to access saved variables after calling backward.

It arises during the second epoch suggesting some part of the graph is being carried over when it shouldn’t be(?). Previous discussions seem to all focus on RNNs and the hidden state needing to be detached from its history. I assume I need to do something similar here but detaching state or output does not solve the error.
If anyone can fix or explain the error that would be greatly appreciated!
Below is the minimal code that produces the error:
def train(model, optimizer, batch_size, n_qubits):
    criterion = nn.MSELoss()

    dataset = DataFactory(batch_size, n_qubits)
    #losses = []
    for epoch in range(5):
        print(epoch)
        x, y = dataset.next_batch() 
        
        state, output = model(x)
        loss = criterion(output.real, y.real)
        
        #Backpropagation
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        #losses.append(loss.item())
        
    return None #losses

n_qubits = 3
batch_size = 1
model = BasicModel(n_qubits, batch_size)
optimizer = torch.optim.AdamW(model.parameters(), lr=0.05)
L = train(model, optimizer, batch_size, n_qubits)

class BasicModel(nn.Module):
    def __init__(self, n_qubits: int, batch_size: int):
        super().__init__()   
        self.n_qubits = n_qubits
        self.batch_size = batch_size
        self.fRx = Rx_layer([*range(n_qubits)], batch_size=batch_size)
        Z = torch.Tensor([[1,0],[0,-1]]).cdouble().reshape((1,1,1,2,2))
        self.Observable = Z.repeat((batch_size, n_qubits, 1, 1, 1))        

    def forward(self, x):
        state = torch.zeros((self.batch_size, self.n_qubits, 1, 2, 1), dtype=torch.cdouble)
        state[:, :, :, 0, 0] = 1 
        Rx1 = Rx_layer([*range(self.n_qubits)], weights=x, batch_size=self.batch_size)
        state = Rx1(state)
        state = self.fRx(state)
        O = torch.matmul(state.transpose(3,4).conj(), torch.matmul(self.Observable, state))
        return state, O.mean(dim=2)[:,:]

class Rx_layer(nn.Module):
    "A layer applying the Rx gate"
    def __init__(self, qubits: list, batch_size: int, weights = None):
        """
        qubits: a list with the index of every qubit this layer is to be applied to
        weights: a tensor of rotation angles, if given from input data
        """
        
        super().__init__()
        
        self.qubits = qubits
        self.num_qubits = len(qubits)
        self.batch_size = batch_size
        if weights is None:
            self.weights = nn.Parameter(torch.Tensor(self.batch_size, self.num_qubits).cdouble())
            nn.init.uniform_(self.weights, 0, 2*np.pi)        
       else:
            if weights.shape[1] == 1 and self.num_qubits > 1:
                self.weights = weights.repeat(1, self.num_qubits)
            elif weights.shape[0] == batch_size and weights.shape[1] == self.num_qubits:
                self.weights = weights
            else:
                raise RuntimeError("Dimensions of weight tensor are incompatable. Check the input has the right batch size and qubit count")
            
        self.U = self.Rx().cdouble()

    def Rx(self):
        a = (self.weights/2).cos().reshape(self.batch_size, self.num_qubits,1,1,1)
        b = (self.weights/2).sin().reshape(self.batch_size, self.num_qubits,1,1,1)
        identity = torch.eye(2).reshape(1,1,1,2,2)
        off_identity = torch.Tensor([[0,1],[1,0]]).reshape(1,1,1,2,2)
        return a*identity - 1j*b*off_identity
        
           
    def forward(self, state):
        """
        Take state to be a tensor with dimension batch x qubits x d&c x 2 x 1
        """
        state[:,self.qubits] = torch.matmul(self.U, state[:,self.qubits]) 

        return state

class DataFactory(Dataset):
    def __init__(self, batch_size: int, input_dim: int):
        self.batch_size = batch_size 
        self.input_dim = input_dim
    
    def next_batch(self):
        X = torch.rand((self.batch_size, self.input_dim)).cdouble()*2 -1
        Y = torch.rand((self.batch_size, self.input_dim)).cdouble()*2 -1
        return X, Y

This is my first post so let me know if there is anything I can clarify (eg I have omitted a mathematical description of the model) or improve!