Flip and transpose with gpu?

#1

Hi, all
My codes for flip and transpose need a lot of time.
Is there any flip method with gpu? or something bad in my code?

def tensor_rot_90(x):
    return x.flip(2).transpose(1, 2)

def apply(func, M):
    tList = [func(m) for m in torch.unbind(M, dim=0)]
    res = torch.stack(tList, dim=0)
    return res

def create_rot_batch(images, rot=['90', '180', '270']):
    rot_batch = images
    rot_labels = torch.ones(len(images)) * 0
    counter = 1

    if '90' in rot:
        images_rot90 = apply(tensor_rot_90_digit, images)
        rot_batch = torch.cat((rot_batch, images_rot90), 0)
        rot_labels = torch.cat((rot_labels, torch.ones(len(images)) * counter), 0)
        counter += 1
         ...