I have a batch sample of size (32, 1024, 3) where 32 is the batch size each of which contains (1024) points with 3 features each.
For each of the sample(among 32) [size (1024,3)] , I am performing few operations which does not  include torch.nn operations(like conv or linear etc) but rather the numpy like operations.
For eg. the operations are like this
curr = 0
batch = torch.randint(0,55,[32,1024,3])
for i in batch:
     for j in i:
          id_ = torch.unique(j)
          id_ = id_[id_!=curr]
         <some other operations>

I wanted to know, can I perform these operations parallely(eliminating the  for i in batch loop)?
The torch.nn APIs automatically perform these tasks on an entire batch without having to explicitly mention it in the code. I wonder if such a similar method can be applied here.
If so, how can I do that?
Tagging @ptrblck for reference!
Thank you!