jarvico:

Now, the thing I need to implement is I want to apply a backward derivative operation on my final output tensor. Actually what I want is to find the derivative of first output scalar in the final output tensor with respect to first input image in the input batch tensor. That is the case for all 64 output scalars and 64 input images.


That would be the usual use case, wouldn’t it? The samples in a batch are independent from each other (besides batch-dependent layers such as batchnorm layers).



 jarvico:

The backward operation there should be applied to every element of final output tensor one by one because for example first output is calculated based on the first input image and x.
and all the final derivatives should be stored in a tensor of shape 64x1x28x28.
(It is composed of sth like: d Out_1/d x_1 , d Out_2/d x_2 , …, d Out_64/d x_64)
If the batch size was 1 I could do it without a problem. But when the batch size is 64 for example, I can’t do it. Can you please help me implement in Pytorch?


If I understand it correctly, your single sample approach would look like this:
# setup
x = torch.randn(64, 2, requires_grad=True)
model = nn.Sequential(
    nn.Linear(2, 10),
    nn.ReLU(),
    nn.Linear(10, 10)
)
    

# single sample approach
for idx, x_ in enumerate(x):
    out = model(x_.unsqueeze(0))
    out.mean(1).backward()
    print(x.grad)
grad_ref = x.grad.clone()

I.e. in each iteration a single sample from the batch would be used and the gradient would be accumulated to the input tensor in the corresponding index. Note how the x.grad attribute is filled during the loop execution.
If the previous approach is what you are targeting, my previous suggestion of using backward(torch.ones_like(out)) should still work:
# clear gradients
x.grad = None

# pass complete batch to model
out = model(x)
out = out.mean(1)
out.backward(torch.ones_like(out))

# compare
print(torch.allclose(grad_ref, x.grad))
> True

As you can see, both approaches yield the same gradients.
If my code snippet is still not mathcing your use case, you would have to provide a code snippet to show what exactly is creating the issues, as I might not understand the terminology clearly.