Solved by John_Smith in post #4 
              
              
                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.r…