ttodic:

Since the transform performs only on numpy arrays, prior to the call I should move the tensor to cpu and convert it to numpy array, right?


Yes, but this would also detach the array from the computation graph and Autograd won’t be able to track these operations (in case you want to use them as differentiable methods). You could either implement the backward pass manually using custom autograd.Functions or check out kornia 2, which provides differentiable transformations (CC @edgarriba as the core dev).



 ttodic:

But I’m not sure how to deal with batch of images, since the network processes whole batch at a time, not just one sample, like the transform works with one sample at a time. The first thing that comes to my mind is to use for loop across all the samples in a batch, could it work this way?


Yes, a loop would work, but could be slow. Since you are already moving the data to the CPU (and thus synchronizing the code), the performance penalty from the loop might not be the bottleneck.



 ttodic:

Also, should I use no_grad() when computing transform? Will this cut the graph or affect training some way?


Yes, it will detach these operations from the computation graph as explained before.