Yes, I did. These are all the cells related to the dataset:
def parse_dataset(dataset):

  dataset.targets = dataset.targets % 2

  return dataset

rescale_data = transforms.Lambda(lambda x : x/255)

# Compose transformations
data_transform = transforms.Compose([
  transforms.Resize(32),
  transforms.RandomHorizontalFlip(),
  transforms.ToTensor(),
  rescale_data,
])

test_transform = transforms.Compose([
  transforms.Resize(32),
  transforms.ToTensor(),
  rescale_data,
])
# Load MNIST dataset with transforms
train_set = torchvision.datasets.MNIST(root=root_dir, train=True, download=True, transform=data_transform)
test_set = torchvision.datasets.MNIST(root=root_dir, train=False, download=True, transform=test_transform)

train_set = parse_dataset(train_set)
test_set = parse_dataset(test_set)

# Dataset len
num_train = len(train_set)
num_test = len(test_set)
print(f"Num. training samples: {num_train}")
print(f"Num. test samples:     {num_test}")

# List of indexes on the training set
train_idx = list(range(num_train))

# List of indexes of the test set
test_idx = list(range(num_test))

# Shuffle the training set

random.shuffle(train_idx)

# Fraction of the original train set that we want to use as validation set
val_frac = 0.1
# Number of samples of the validation set
num_val = int(num_train * val_frac)
num_train = num_train - num_val

# Split training set
val_idx = train_idx[num_train:]
train_idx = train_idx[:num_train]

print(f"{num_train} samples used as train set")
print(f"{num_val}  samples used as val set")

len(train_idx)

val_set = Subset(train_set, val_idx)
train_set = Subset(train_set, train_idx)

# Let's split in loader_a and loader_b
num_train = int(len(train_set)/2)
num_test = int(len(test_set)/2)

# List of indexes on the training set
train_idx = list(range(num_train*2))

# List of indexes of the test set
test_idx = list(range(num_test*2))

# Split training set
train_a_idx = train_idx[num_train:]
train_b_idx = train_idx[:num_train]

# Split test set
test_a_idx = test_idx[num_test:]
test_b_idx = test_idx[:num_test]

train_set_a1 = Subset(train_set, train_a_idx)
train_set_b1 = Subset(train_set, train_b_idx)

test_set_a = Subset(test_set, test_a_idx)
test_set_b = Subset(test_set, test_b_idx)

print(f"{num_train} samples used as train set")
print(f"{num_test}  samples used as test set")

num_train_a1 = len(train_set_a1)
num_train_b1 = len(train_set_b1)
# Fraction of the original train set that we want to use as validation set
val_frac = 0.1
# Number of samples of the validation set
num_val_a = int(num_train_a1 * val_frac)
num_train_a = num_train_a1 - num_val_a

num_val_b = int(num_train_b1 * val_frac)
num_train_b = num_train_b1 - num_val_b

# Split training set
val_idx_a = train_idx[num_train_a:]
train_idx_a = train_idx[:num_train_a]
val_set_a = Subset(train_set_a1, val_idx_a)
train_set_a = Subset(train_set_a1, train_idx_a)

val_idx_b = train_idx[num_train_b:]
train_idx_b = train_idx[:num_train_b]
val_set_b = Subset(train_set_b1, val_idx_b)
train_set_b = Subset(train_set_b1, train_idx_b)

print(f"{num_train_a} samples used as train set a")
print(f"{num_val_a}  samples used as val set a")
print(f"{num_train_b} samples used as train set b")
print(f"{num_val_b}  samples used as val set b")