I’ve tried a simple example and it seems to me it should work.
import torch
import torch.nn as nn

# create 4 tensorts to simulate 4 images with size 2*2
a,b,c,d = torch.arange(16).reshape(-1,2).chunk(4)

# set them to require grads
a.requires_grad=True
b.requires_grad=True
c.requires_grad=True
d.requires_grad=True

# concat 4 imgs into 1 img with size 4*4
e = torch.cat((torch.cat((a,b),1), torch.cat((c,d),1)))

# simulate the forward pass with a liniear transformation
f = e.mm(torch.arange(8).reshape(4,2))

# calculate loss using cross entropy
loss = nn.CrossEntropyLoss()(f, torch.tensor([0,1,0,0]))

# backword loss
loss.backward()

# check grads of a w.r.t loss
print(a.grad)