Hi @ptrblck,
I really appreciate your effort of writing up a code to explain what is happening. And yeah, your code captures what I was looking for.
I extended your code from one layer to 2 (since weights don’t affect gradients in a single layer net) to actually see the effect of changing weights after loss computation.
According to the observation, it looks like the gradients are computed during ‘backward()’ call using the new weights i.e. W2.
import torch
from torch.autograd import Variable
from torch import nn
from torch import optim

from graphviz import Digraph
# make_dot was moved to https://github.com/szagoruyko/pytorchviz
from torchviz import make_dot


w0 = Variable(torch.ones(2, 2), requires_grad=True) 
w = Variable(torch.ones(2, 1), requires_grad=True)

x = Variable(torch.ones(1, 2))
y = Variable(torch.ones(1, 1))

criterion = nn.MSELoss()
optimizer = optim.SGD([w0,w], lr=0.01)

for i in range(1):

    output = x.mm(w0).mm(w)

    # if i == 0:
    #     make_dot(output).view()

    optimizer.zero_grad()
    loss = criterion(output, y)
    # w0.data.fill_(2) # comment / uncomment
    # w.data.fill_(2) # comment / uncomment
    loss.backward()

    print('w0 {}grad {}'.format(w0.data, w0.grad.data))
    print('w {}grad {}'.format(w.data, w.grad.data))


    optimizer.step()
    print('w0 {}grad {}'.format(w0.data, w0.grad.data))
    print('w {}grad {}'.format(w.data, w.grad.data))
    optimizer.zero_grad()