Solved by ptrblck in post #8 
              
              
                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). 

If I understand it correctly, your single sample approach would look like this: 
# setup
x = torch.randn(64, 2, requires_grad=True)
mode…