def get_subset(options, dataset, class_a, class_b):
    """
    Returns a subset of the dataset that only contains the
    classes (class_a) and (class_b)
    """
    x_train, y_train = dataset
    train_idx = (y_train == class_a) ^ (y_train == class_b)

    subset_train_x = x_train[train_idx]
    subset_train_y = y_train[train_idx]

    # relabel to +/-1
    subset_train_y[subset_train_y == class_a] = -1
    subset_train_y[subset_train_y == class_b] = 1

    return subset_train_x, subset_train_y