I’m attempting to implement a custom single forward pass training algorithm using PyTorch. Since I don’t require backpropagation, I am manually updating the weights of the neural network. However, I can’t seem to get this to work correctly.
After the first pass, I repeatedly get the error that I’m trying to backward through the computational graph for a second time, despite having zeroed out the gradients in the model. Not sure where I’m going wrong.

class OneD_NN_LQR:
    # Initially we will ignore batch size
    def __init__(self, hidden_units, learning_rate_param_C=0.05, batch_size=100):
        # Single layer neural network for the control, of the form f(x) = sum(c_i * g(w_i * x + b_i))
        # We will use a sigmoid activation function (i.e. g = sigmoid)
        self.C = learning_rate_param_C
        self.N = batch_size
        self.hidden_units = hidden_units
        self.dim = 1

        self.layer1 = torch.nn.Linear(in_features=self.dim, out_features=self.hidden_units)
        self.activation = torch.nn.ReLU()
        self.layer2 = torch.nn.Linear(in_features=self.hidden_units, out_features=self.dim, bias=False)
        self.model = torch.nn.Sequential(
            self.layer1,
            self.activation,
            self.layer2
        )
        self.w = self.layer1.weight
        self.b = self.layer1.bias
        self.c = self.layer2.weight

        self.Xtilde_w = torch.zeros((self.hidden_units,)).unsqueeze(1)
        self.Xtilde_c = torch.zeros((self.hidden_units,)).unsqueeze(1)
        self.Xtilde_b = torch.zeros((self.hidden_units,))

        self.X = torch.ones((self.dim,), requires_grad=True)
        self.f_x = self.forward(self.X)

        #self.grads = torch.autograd.grad(self.f_x, inputs=[self.X, self.layer2.weight, self.layer1.bias, self.layer1.weight])
        self.f_x.backward()
        # self.grad_x = self.grads[0]
        # self.grad_c = self.grads[1].T
        # self.grad_b = self.grads[2]
        # self.grad_w = self.grads[3]
        self.grad_x = self.X.grad
        self.grad_c = self.c.grad.T
        self.grad_b = self.b.grad
        self.grad_w = self.w.grad
        
        self.time = 0
        
    def step(self, delta):
        # Stepping also involves updating the values of f(x) and f'(x)
        self.time += delta

        self.step_X(delta)
        self.step_Xtilde(delta)
        self.step_theta(delta)
        
        self.model.zero_grad()

        self.f_x = self.model.forward(self.X)

        print(self.f_x)
        #self.grads = torch.autograd.grad(self.f_x, inputs=[self.X, self.layer2.weight, self.layer1.bias, self.layer1.weight])
        self.f_x.backward()
        # self.grad_x = self.grads[0]
        # self.grad_c = self.grads[1].T
        # self.grad_b = self.grads[2]
        # self.grad_w = self.grads[3]
        self.grad_x = self.X.grad
        self.grad_c = self.c.grad.T
        self.grad_b = self.b.grad
        self.grad_w = self.w.grad
        
        return self.w, self.c, self.b

    def step_theta(self, delta):
        next_dw, next_dc, next_db = self.next_dtheta(delta)

        with torch.no_grad():
            self.layer1.weight.sub_(next_dw)
            self.layer1.bias.sub_(next_db)
            self.layer2.weight.sub_(next_dc.T)
            self.model.zero_grad()

    def step_X(self, delta):
        next_dX = self.next_dX(delta)
        self.X = self.X + next_dX

    def step_Xtilde(self, delta):
        next_dXtilde_w, next_dXtilde_c, next_dXtilde_b = self.next_dXtilde(delta)

        self.Xtilde_w = self.Xtilde_w + next_dXtilde_w
        self.Xtilde_c = self.Xtilde_c + next_dXtilde_c
        self.Xtilde_b = self.Xtilde_b + next_dXtilde_b
        

    def next_dtheta(self, delta):
        alpha = self.get_learning_rate(self.C, self.time)

        dw = alpha * (2 * self.X * self.Xtilde_w + 2 * self.f_x * (self.grad_w + self.grad_x * self.Xtilde_w)) * delta
        db = alpha * (2 * self.X * self.Xtilde_b + 2 * self.f_x * (self.grad_b + self.grad_x * self.Xtilde_b)) * delta
        dc = alpha * (2 * self.X * self.Xtilde_w + 2 * self.f_x * (self.grad_c + self.grad_x * self.Xtilde_c)) * delta
        
        return dw, dc, db

    def get_learning_rate(self, c, time):
        if time > 500: return c / 10
        if time > 100: return c / 5
        if time > 50: return c / 2
        return c

    def next_dXtilde(self, delta):
        
        dXtilde_w = (- self.Xtilde_w + self.grad_w + self.grad_x * self.Xtilde_w) * delta
        dXtilde_b = (- self.Xtilde_b + self.grad_b + self.grad_x * self.Xtilde_b) * delta
        dXtilde_c = (- self.Xtilde_c + self.grad_c + self.grad_x * self.Xtilde_c) * delta
        
        return dXtilde_w, dXtilde_c, dXtilde_b

    def next_dX(self, delta):
        to_return = (-self.X + self.f_x) * delta + torch.normal(0, 1, size=(self.dim,)) * (delta ** 0.5)
        return to_return

    def forward(self, x):
        #to_return = torch.unsqueeze(torch.sum(self.c * self.activation(self.w * x + self.b), axis=1), 1)
        to_return = self.model.forward(x)
        return to_return