I wrote a short snippet to test ‘_stateless’. In my example I tried to update weights of a network in an inner optimization loop and to learn the learning rate of the weight updates in an outer optimization loop (meta-optimization). I’m getting the error:

RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation: [torch.FloatTensor [3, 10]], which is output 0 of AsStridedBackward0, is at version 12; expected version 2 instead. Hint: enable anomaly detection to find the operation that failed to compute its gradient, with torch.autograd.set_detect_anomaly(True).

My code snippet is:
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 self.fc2(y)

epochs = 20
N = 100
dataset = DataLoader(dataset=MyDataset(N), batch_size=10)
model = MyModel()
loss_func = nn.CrossEntropyLoss()

optim = optim.Adam([model.alpha], lr=1e-3)

params = dict(model.named_parameters())
for i in range(epochs):
    model.train()
    train_loss = 0
    for batch_idx, (x, y) in enumerate(dataset):
        logits = _stateless.functional_call(model, params, x)             # predict
        loss_inner = loss_func(logits, y)                                 # loss
        optim.zero_grad()                                                 # reset grad
        loss_inner.backward(create_graph=True, inputs=params.values())    # compute grad
        train_loss += loss_inner.item()                                   # store loss
        for k, p in params.items():
            if k is not 'alpha' and k is not 'beta':
                p.update = - model.alpha * p.grad
                params[k] = p + p.update                      # update weight

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

From the error message, I understand that the issue comes from weight update for the weights of the second layer of the network, which points to an error in my inner loop optimization. Do you have suggestions? Thanks!
PS: I thought this is not directly relevant to the feature request, but I can move my question to the github issue if you prefer to discuss there. Thank you!