def train_net(train, val, model, name):
transformations_train = transforms.apply_chain([
transforms.to_float,
transforms.augment_color(0.5),
transforms.random_fliplr(),
transforms.random_flipud(),
transforms.augment(),
torchvision.transforms.ToTensor()
])
transformations_val = transforms.apply_chain([
    transforms.to_float,
    torchvision.transforms.ToTensor()
])

dset_train = HumanProtienJPGDataset(train, paths.train_jpg, transformations_train, divide=False)
train_loader = DataLoader(dset_train,
                          batch_size=2,
                          shuffle=True,
                          #num_workers=10,
                          pin_memory=True)

dset_val = HumanProtienJPGDataset(val, paths.train_jpg, transformations_val, divide=False)
val_loader = DataLoader(dset_val,
                        batch_size=32,
                        #num_workers=10,
                        pin_memory=True)

If i reduce the batch size to 1 i get the following error
RuntimeError: Expected 4-dimensional input for 4-dimensional weight [64, 3, 7, 7], but got 3-dimensional input of size [3, 512, 512] instead