I was trying to add a Hebbian learning term (based on which weight update is proportional to the product of pre and post-synaptic activations) to the inner optimization. To that end, I modified the code as following:
import torch
from torch import nn, optim
from torch.utils.data import Dataset, DataLoader

from torch.nn.utils import _stateless


class MyDataset(Dataset):
    def __init__(self, N):
        self.N = N
        self.x = torch.rand(self.N, 10)
        self.y = torch.randint(0, 3, (self.N,))

    def __len__(self):
        return self.N

    def __getitem__(self, idx):
        return self.x[idx], self.y[idx]


class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.fc1 = nn.Linear(10, 10)
        self.fc2 = nn.Linear(10, 3)

        self.relu = nn.ReLU()

        self.alpha = nn.Parameter(torch.randn(1))
        self.beta = nn.Parameter(torch.randn(1))

    def forward(self, x):
        y = self.relu(self.fc1(x))
        return (x, y), self.fc2(y)

def Optim(params, alpha, beta, y, logits):

    # -- add network output to activations
    softmax = nn.Softmax(dim=1)
    activations = []
    for item in y:
        activations.append(item)
    activations.append(softmax(logits))

    i = 0
    for k, p in params.items():
        print(k[4:])
        if k[4:] == 'weight':
            p.update = - alpha * p.grad + beta * torch.matmul(activations[i+1].T, activations[i])
            params[k] = p + p.update   # update weight
        elif k[4:] == 'bias':
            p.update = - alpha * p.grad + beta * activations[i + 1].squeeze(0)
            params[k] = p + p.update   # update weight
            i += 1

epochs = 20
N = 10
dataset = DataLoader(dataset=MyDataset(N), batch_size=1)
model = MyModel()
loss_func = nn.CrossEntropyLoss()
optim = optim.Adam([model.alpha], lr=1e-3)
torch.autograd.set_detect_anomaly(True)

for i in range(epochs):
    model.train()
    train_loss = 0
    params = dict(model.named_parameters())
    for batch_idx, (x, y) in enumerate(dataset):
        params = {k: v.clone() for k, v in params.items()}
        activations, logits = _stateless.functional_call(model, params, x)             # predict
        loss_inner = loss_func(logits, y)                                 # loss
        loss_inner.backward(create_graph=True, inputs=params.values())    # compute grad
        train_loss += loss_inner.item()                                   # store loss
        Optim(params, model.alpha, model.beta, activations, logits)

    print('Train Epoch: {}\tLoss: {:.6f}'.format(i, train_loss / N))
    activations, logits = _stateless.functional_call(model, params, x)                 # predict
    loss_meta = loss_func(logits, y)
    loss_meta.backward()
    optim.step()

This gives me the error

RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation: [torch.FloatTensor [1]] is at version 11; expected version 10 instead. Hint: the backtrace further above shows the operation that failed to compute its gradient. The variable in question was changed in there or anywhere later. Good luck!

when I change the line
    Optim(params, model.alpha, model.beta, activations, logits)

to
    Optim(params, model.alpha, model.beta.clone(), activations, logits)

the problem vanishes. Since model.alpha does not need cloning,

I’m wondering if using model.beta.clone() is the correct solution?
Why unlike model.beta, cloning is not needed when passing model.alpha to Optim?