IT was something to do with
transforms.RandomErasing(p=0.4, scale=(0.09, 0.25), ratio=(0.3, 3.3), value=0, inplace=False),

As this works fine
train_transform = transforms.Compose([
    transforms.Resize(size),
    transforms.RandomHorizontalFlip(),
    transforms.RandomRotation(10),
    transforms.ToTensor(),
    transforms.Normalize(mean, std)
])

val_transform = transforms.Compose([
    transforms.Resize(size),
    transforms.ToTensor(),
    transforms.Normalize(mean, std)
])

I have seen on stack overflow that this can cause some issues and a way to work with it is to go into the source code and change size to shape to get it to work. I dont feel comfertable with that so I will pass.