Hi I am trying to create a custom optimizer that simulates the non-linear and asymmetric behavior of Resistive RAM. That is, simply speaking, when the weight is updated, change in weight is a function of the current value of weight.
Below is my code for custom optimizer. It bases SGD source code so the only part that I changed is def step.
'''Custom Optimizer'''

from torch.optim.optimizer import Optimizer, required
import copy

class RPU_SGD(Optimizer):
    def __init__(self, params, lr = required, momentum =0, dampening =0, weight_decay=0, nesterov=False):
        if lr is not required and lr < 0.0:
            raise ValueError("Invalid learning rate: {}".format(lr))
        if momentum < 0.0:
            raise ValueError("Invalid momentum value: {}".format(momentum))
        if weight_decay<0.0:
            raise ValueError("Invalid weight_decay value: {}".format(weight_decay))
        
        defaults = dict(lr=lr, momentum=momentum, dampening=dampening, weight_decay=weight_decay, nesterov=nesterov)
        super(RPU_SGD, self).__init__(params, defaults)

    def __setstate__(self, state):
        super(RPU_SGD, self).__setstate__(state)
        
    def step(self, closure=None):        
        loss =None
        if closure is not None:
            with torch.enable_grad():
                loss = closure()
        
        for group in self.param_groups:
            weight_decay = group['weight_decay']
            momentum = group['momentum']
            dampening = group['dampening']
            nesterov = group['nesterov']
        
            for p in group['params']:
                delta_w_0 = 0.012
                slope_p = 9.63
                slope_n = 9.63
    
                if p.grad is None:
                    continue
                d_p=p.grad
                
                if len(p.size())==2:
                    for i in range(p.size()[0]):
                        for j in range(p.size()[1]):
                            delta_w_pos = delta_w_0 * (1 - slope_p * p[i][j])
                            delta_w_neg = delta_w_0 * (1 + slope_n * p[i][j])

                            if d_p[i][j]>0:
                                gradient = delta_w_pos
                                p[i][j].add_(gradient, alpha=-group['lr'])

                            elif d_p[i][j]<0:
                                gradient = delta_w_neg
                                p[i][j].add_(gradient, alpha=-group['lr'])

                            elif d_p[i][j]==0:
                                graident =0
                                p[i][j].add_(gradient, alpha=-group['lr'])

   
        return loss

And below is my training code.
'''Training the Network'''
net.train()

for epoch in range(epochs):
#     net.train()
    
    running_loss = 0
    for i,data in enumerate(train_loader, 0):
        inputs,labels = data
        
        optimizer.zero_grad()
        pred = net(inputs)
        loss = criterion(pred, labels)
        loss.backward()
        optimizer.step()
        
        running_loss+=loss.item()
        if (i+1)%100==0:
            print(f"epoch: {epoch}/{epochs} | step: {i+1}/{len(train_loader)}, loss: {running_loss/100:.4f}")
            running_loss = 0
        
        new_params = list(net.parameters())

print("End of Training")


When I run them I get error of:
RuntimeError                              Traceback (most recent call last)
<ipython-input-64-b27ef8dbc022> in <module>
     12         pred = net(inputs)
     13         loss = criterion(pred, labels)
---> 14         loss.backward()
     15 #         import pdb; pdb.set_trace()
     16 #         with torch.no_grad():

~/anaconda3/lib/python3.8/site-packages/torch/tensor.py in backward(self, gradient, retain_graph, create_graph)
    193                 products. Defaults to ``False``.
    194         """
--> 195         torch.autograd.backward(self, gradient, retain_graph, create_graph)
    196 
    197     def register_hook(self, hook):

~/anaconda3/lib/python3.8/site-packages/torch/autograd/__init__.py in backward(tensors, grad_tensors, retain_graph, create_graph, grad_variables)
     95         retain_graph = create_graph
     96 
---> 97     Variable._execution_engine.run_backward(
     98         tensors, grad_tensors, retain_graph, create_graph,
     99         allow_unreachable=True)  # allow_unreachable flag

RuntimeError: leaf variable has been moved into the graph interior

I have two issues:

I am not sure why leaf variable has been moved into the graph interior(tbh idk what it means too).
2.I used for loop to update each weight parameters, yet this, i reckon, would be very time-consuming. Would there be any better way to do this?

Thanks in advance:)