I have an image, and I create four rotated versions of it using.
def rotate4dtensor(x):
    x90 = x.transpose(2, 3)
    x180 = x.flip(2)
    x270 = x.transpose(2, 3).flip(3)
    return [x,x90,x180,x270]

and now when during training I call the function in this manner
target = target.cuda()
        input = input.cuda()
        input = torch.autograd.Variable(input)
        target_var = torch.autograd.Variable(target)
        output = model(input)
        loss = criterion(output, target_var)
        
        
        
        #####vectorized version
        losses_rot = []
        #rotation loss for supervised
        sup_rot = rotate4dtensor(input); tgts = []
        sup_rot4 = torch.stack(sup_rot)
        out = model(sup_rot4,is_rotation=True)

It throws the following error when no of gpus 3,5,6 but works fine with 1,2,and 4 why?
image (2)3160×916 797 KB