Currently I am trying like the following:
def mytransform(self, x):
     for i in x.shape[0]:
          x[i,:,:,:] = self.transform(x[i,:,:,:])
     return x

def train(self):
    for (x,y) in loader:
         xt = self.mytransform(x)
         process(xt)
         ...

but is there any efficient way to do it?