ValueError: Expected input batch_size (384) to match target batch_size (128).

torch.manual_seed(43)

val_size = 5000

train_size = len(dataset) - val_size