Repository: apple/pfl-research (language: python)

Issue:
rdar://120413769 CI for CIFAR10 code
CI for running cifar10 code with tiny dataset. only takes 10s to run training.



Discussion / hints:

Fix patch:
diff --git a/.circleci/config.yml b/.circleci/config.yml
index f1e481d..409c55a 100644
--- a/.circleci/config.yml
+++ b/.circleci/config.yml
@@ -23,6 +23,16 @@ commands:
           command: |
             pip install poetry
             poetry install --only dev -q || true
+  benchmark_install_steps:
+    steps:
+      - checkout
+      - run:
+          name: Install dependencies
+          command: |
+            cd benchmarks
+            pip install poetry
+            poetry install -E tf -E pytorch --no-root || true
+            echo 'export PYTHONPATH=.' >> $BASH_ENV
 
 jobs:
   code-quality:
@@ -68,6 +78,33 @@ jobs:
           name: "Test with PyTorch"
           command: poetry run tox -e 'py310-test-pytorch'
 
+  check-run-cifar10:
+    executor: linux-python
+    steps:
+      - benchmark_install_steps
+      - run:
+          name: "Download data"
+          command: |
+            mkdir -p ./benchmarks/data/cifar10
+            curl https://pfl-data.s3.us-east-2.amazonaws.com/cifar10_ci/cifar10_train.p \
+              -o ./benchmarks/data/cifar10/cifar10_train.p
+            curl https://pfl-data.s3.us-east-2.amazonaws.com/cifar10_ci/cifar10_test.p \
+              -o ./benchmarks/data/cifar10/cifar10_test.p
+      - run:
+          name: "(PyTorch) run a few iterations of training"
+          command: |
+            cd benchmarks
+            poetry run python image_classification/pytorch/train.py \
+              --args_config image_classification/configs/baseline.yaml \
+              --central_num_iterations 11 --cohort_size 2 --val_cohort_size 2
+      - run:
+          name: "(TF) run a few iterations of training"
+          command: |
+            cd benchmarks
+            poetry run python image_classification/tf/train.py \
+              --args_config image_classification/configs/baseline.yaml \
+              --central_num_iterations 11 --cohort_size 2 --val_cohort_size 2
+
   publish-wheel:
     executor: linux-python
     steps:
@@ -94,6 +131,7 @@ workflows:
       - build-documentation-wheel
       - test-tf
       - test-pytorch
+      - check-run-cifar10
 
       # TODO: rdar://120414177 (optional CI for docker image (build only))
       #- build-images:
@@ -118,3 +156,5 @@ workflows:
 
 
 
+
+
diff --git a/.gitignore b/.gitignore
index 9636247..897e13c 100644
--- a/.gitignore
+++ b/.gitignore
@@ -167,3 +167,5 @@ docs/build
 .DS_Store
 
 benchmarks/data
+benchmarks/checkpoints/
+benchmarks/metrics.csv
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index a832fa8..b5c3678 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -36,7 +36,6 @@ repos:
       - id: yapf
         args:
           - --in-place
-        files: ^(pfl|tests)/
   # TODO: fix ruff errors and remove pylint, then enable this.
   # Style checks, common errors, etc.
   #- repo: https://github.com/charliermarsh/ruff-pre-commit
diff --git a/benchmarks/dataset/argument_parsing.py b/benchmarks/dataset/argument_parsing.py
index fa91acf..0deddc5 100644
--- a/benchmarks/dataset/argument_parsing.py
+++ b/benchmarks/dataset/argument_parsing.py
@@ -18,21 +18,20 @@ def add_dataset_arguments(
     arguments depending on the dataset specified in ``dataset``.
     """
 
-    parser.add_argument(
-        '--dataset',
-        choices=[
-            'cifar10', 'cifar10_iid', 'femnist', 'femnist_digits', 'reddit',
-            'flair', 'flair_iid', 'flair_pytorch', 'stackoverflow'
-        ],
-        default='cifar10',
-        help='Which dataset to train on')
-
-    parser.add_argument(
-        '--minimum_num_datapoints_per_user',
-        type=int,
-        default=1,
-        help='The minimum number of samples that allows a'
-        'simulated user to participate.')
+    parser.add_argument('--dataset',
+                        choices=[
+                            'cifar10', 'cifar10_iid', 'femnist',
+                            'femnist_digits', 'reddit', 'flair', 'flair_iid',
+                            'flair_pytorch', 'stackoverflow'
+                        ],
+                        default='cifar10',
+                        help='Which dataset to train on')
+
+    parser.add_argument('--minimum_num_datapoints_per_user',
+                        type=int,
+                        default=1,
+                        help='The minimum number of samples that allows a'
+                        'simulated user to participate.')
 
     # Get the value of `dataset` argument and dynamically add
     # arguments depending on which dataset is chosen.
@@ -51,11 +50,10 @@ def add_dataset_arguments(
             help='Alpha for Dir(alpha) non-iid class partitioning, '
             'not used for cifar10_iid')
     elif known_args.dataset in {'reddit', 'stackoverflow'}:
-        parser.add_argument(
-            "--data_fraction",
-            type=float,
-            default=1.0,
-            help='Fraction of total dataset to use.')
+        parser.add_argument("--data_fraction",
+                            type=float,
+                            default=1.0,
+                            help='Fraction of total dataset to use.')
 
         parser.add_argument(
             "--central_data_fraction",
@@ -78,17 +76,15 @@ def add_dataset_arguments(
     elif known_args.dataset in ['flair', 'flair_iid', 'flair_pytorch']:
         parser = add_artificial_fed_dataset_arguments(parser)
 
-        parser.add_argument(
-            '--use_fine_grained_labels',
-            action=store_bool,
-            default=False,
-            help='Whether to use fine-grained label taxonomy.')
+        parser.add_argument('--use_fine_grained_labels',
+                            action=store_bool,
+                            default=False,
+                            help='Whether to use fine-grained label taxonomy.')
 
-        parser.add_argument(
-            '--max_num_user_images',
-            type=int,
-            default=100,
-            help='Maximum number of images per user')
+        parser.add_argument('--max_num_user_images',
+                            type=int,
+                            default=100,
+                            help='Maximum number of images per user')
 
     return parser
 
@@ -100,19 +96,17 @@ def add_artificial_fed_dataset_arguments(
     ``ArtificialFederatedDataset``, e.g. for CIFAR10.
     """
 
-    parser.add_argument(
-        '--datapoints_per_user_distribution',
-        choices=['constant', 'poisson'],
-        default='constant',
-        help='Distribution of the number of samples that '
-        'simulated users have to train on.')
+    parser.add_argument('--datapoints_per_user_distribution',
+                        choices=['constant', 'poisson'],
+                        default='constant',
+                        help='Distribution of the number of samples that '
+                        'simulated users have to train on.')
 
-    parser.add_argument(
-        '--mean_datapoints_per_user',
-        type=float,
-        default=50,
-        help='Mean parameter for the number of samples that '
-        'the simulated users have to train on.')
+    parser.add_argument('--mean_datapoints_per_user',
+                        type=float,
+                        default=50,
+                        help='Mean parameter for the number of samples that '
+                        'the simulated users have to train on.')
 
     return parser
 
@@ -155,9 +149,10 @@ def draw_truncated_poisson():
     return draw_num_datapoints_per_user
 
 
-def get_datasets(args: argparse.Namespace
-                 ) -> Tuple[FederatedDatasetBase, FederatedDatasetBase,
-                            Dataset, Dict[str, Any]]:
+def get_datasets(
+    args: argparse.Namespace
+) -> Tuple[FederatedDatasetBase, FederatedDatasetBase, Dataset, Dict[str,
+                                                                     Any]]:
     """
     Create a federated dataset for training, a federated dataset for evalution
     and a central dataset for central evaluation.
@@ -186,9 +181,10 @@ def get_datasets(args: argparse.Namespace
             args.minimum_num_datapoints_per_user)
 
         if args.dataset == 'cifar10':
-            datasets = cifar10.make_cifar10_datasets(
-                args.data_path, user_dataset_len_sampler, numpy_to_tensor,
-                args.partition_alpha)
+            datasets = cifar10.make_cifar10_datasets(args.data_path,
+                                                     user_dataset_len_sampler,
+                                                     numpy_to_tensor,
+                                                     args.partition_alpha)
         elif args.dataset == 'cifar10_iid':
             datasets = cifar10.make_cifar10_iid_datasets(
                 args.data_path, user_dataset_len_sampler, numpy_to_tensor)
diff --git a/benchmarks/dataset/cifar10/download_preprocess.py b/benchmarks/dataset/cifar10/download_preprocess.py
index 448f17d..96059c3 100644
--- a/benchmarks/dataset/cifar10/download_preprocess.py
+++ b/benchmarks/dataset/cifar10/download_preprocess.py
@@ -38,8 +38,8 @@ def dl_preprocess_and_dump(output_dir: str):
         print(f'Found {tar_path} on disk, skip download')
 
     # Extract tar file.
-    subprocess.run(
-        f"tar -zxf {tar_path} -C {raw_output_dir}".split(), check=True)
+    subprocess.run(f"tar -zxf {tar_path} -C {raw_output_dir}".split(),
+                   check=True)
 
     # Merge all files into a pickle with train data and a pickle with test data.
     def merge_data_and_dump(data_paths, output_file_name):
@@ -60,10 +60,10 @@ def merge_data_and_dump(data_paths, output_file_name):
             pickle.dump((images, labels), f)
         print(f'Saved processed data to {out_file_path}')
 
-    merge_data_and_dump(
-        glob(raw_output_dir + '/**/data_batch*'), 'cifar10_train.p')
-    merge_data_and_dump(
-        glob(raw_output_dir + '/**/test_batch*'), 'cifar10_test.p')
+    merge_data_and_dump(glob(raw_output_dir + '/**/data_batch*'),
+                        'cifar10_train.p')
+    merge_data_and_dump(glob(raw_output_dir + '/**/test_batch*'),
+                        'cifar10_test.p')
 
 
 if __name__ == '__main__':
diff --git a/benchmarks/dataset/cifar10/numpy.py b/benchmarks/dataset/cifar10/numpy.py
index 788c623..51e2f38 100644
--- a/benchmarks/dataset/cifar10/numpy.py
+++ b/benchmarks/dataset/cifar10/numpy.py
@@ -75,8 +75,9 @@ def make_iid_federated_dataset(
     images = numpy_to_tensor(images)
     labels = numpy_to_tensor(labels)
 
-    return ArtificialFederatedDataset.from_slices(
-        [images, labels], data_sampler, user_dataset_len_sampler)
+    return ArtificialFederatedDataset.from_slices([images, labels],
+                                                  data_sampler,
+                                                  user_dataset_len_sampler)
 
 
 def make_central_dataset(images: np.ndarray, labels: np.ndarray) -> Dataset:
@@ -89,10 +90,10 @@ def make_central_dataset(images: np.ndarray, labels: np.ndarray) -> Dataset:
 
 
 def make_cifar10_datasets(
-        data_dir: str,
-        user_dataset_len_sampler: Callable,
-        numpy_to_tensor: Callable,
-        alpha: float = 0.1,
+    data_dir: str,
+    user_dataset_len_sampler: Callable,
+    numpy_to_tensor: Callable,
+    alpha: float = 0.1,
 ) -> Tuple[FederatedDataset, FederatedDataset, Dataset, Dict[str, Any]]:
     """
     Create a train and val ``ArtificialFederatedDataset`` as well as a
@@ -121,8 +122,8 @@ def make_cifar10_datasets(
 
 
 def make_cifar10_iid_datasets(
-        data_dir: str, user_dataset_len_sampler: Callable,
-        numpy_to_tensor: Callable
+    data_dir: str, user_dataset_len_sampler: Callable,
+    numpy_to_tensor: Callable
 ) -> Tuple[ArtificialFederatedDataset, ArtificialFederatedDataset, Dataset,
            Dict[str, Any]]:
     """
diff --git a/benchmarks/dataset/femnist.py b/benchmarks/dataset/femnist.py
index aeda553..26f1a39 100644
--- a/benchmarks/dataset/femnist.py
+++ b/benchmarks/dataset/femnist.py
@@ -12,10 +12,11 @@
 from pfl.data.dataset import Dataset
 
 
-def _load_h5_into_dict(h5_file_path: str,
-                       digits_only: bool,
-                       numpy_to_tensor: Callable = lambda x: x
-                       ) -> Dict[str, List[np.ndarray]]:
+def _load_h5_into_dict(
+        h5_file_path: str,
+        digits_only: bool,
+        numpy_to_tensor: Callable = lambda x: x
+) -> Dict[str, List[np.ndarray]]:
     """
     Load data into memory and create a mapping from user ids to that
     user's data.
@@ -89,9 +90,9 @@ def make_central_dataset(h5_file_path: str,
 
 
 def make_femnist_datasets(
-        data_dir: str,
-        digits_only: bool = False,
-        numpy_to_tensor: Callable = lambda x: x,
+    data_dir: str,
+    digits_only: bool = False,
+    numpy_to_tensor: Callable = lambda x: x,
 ) -> Tuple[FederatedDataset, FederatedDataset, Dataset, Dict[str, Any]]:
     """
     Create a train and val ``FederatedDataset`` as well as a central dataset
@@ -105,8 +106,9 @@ def make_femnist_datasets(
     # data
     training_federated_dataset = make_federated_dataset(
         train_h5_file_path, digits_only, numpy_to_tensor)
-    val_federated_dataset = make_federated_dataset(
-        val_h5_file_path, digits_only, numpy_to_tensor)
+    val_federated_dataset = make_federated_dataset(val_h5_file_path,
+                                                   digits_only,
+                                                   numpy_to_tensor)
     central_data = make_central_dataset(val_h5_file_path, digits_only)
 
     return training_federated_dataset, val_federated_dataset, central_data, {}
diff --git a/benchmarks/dataset/flair/__init__.py b/benchmarks/dataset/flair/__init__.py
index 6bc661a..e09016c 100644
--- a/benchmarks/dataset/flair/__init__.py
+++ b/benchmarks/dataset/flair/__init__.py
@@ -38,9 +38,10 @@ def make_flair_datasets(data_path: str, use_fine_grained_labels: bool,
     training_federated_dataset = make_federated_dataset(
         data_path, 'train', use_fine_grained_labels, max_num_user_images,
         numpy_to_tensor)
-    val_federated_dataset = make_federated_dataset(
-        data_path, 'val', use_fine_grained_labels, max_num_user_images,
-        numpy_to_tensor)
+    val_federated_dataset = make_federated_dataset(data_path, 'val',
+                                                   use_fine_grained_labels,
+                                                   max_num_user_images,
+                                                   numpy_to_tensor)
 
     central_data, metadata = get_central_data_and_metadata(
         data_path, use_fine_grained_labels)
@@ -80,8 +81,9 @@ def make_flair_pytorch_datasets(data_path: str, use_fine_grained_labels: bool,
 
     training_federated_dataset = make_federated_dataset(
         data_path, 'train', use_fine_grained_labels, max_num_user_images)
-    val_federated_dataset = make_federated_dataset(
-        data_path, 'val', use_fine_grained_labels, max_num_user_images)
+    val_federated_dataset = make_federated_dataset(data_path, 'val',
+                                                   use_fine_grained_labels,
+                                                   max_num_user_images)
 
     central_data, metadata = get_central_data_and_metadata(
         data_path, use_fine_grained_labels)
diff --git a/benchmarks/dataset/flair/common.py b/benchmarks/dataset/flair/common.py
index d5d1c8f..4a0ae32 100644
--- a/benchmarks/dataset/flair/common.py
+++ b/benchmarks/dataset/flair/common.py
@@ -43,10 +43,10 @@ def get_label_mapping(hdf5_path: str,
             return json.loads(h5['/metadata/label_mapping'][()])
 
 
-def get_training_channel_mean_stddevs(hdf5_path: str,
-                                      data_fraction: float = 0.1,
-                                      seed: Optional[int] = None
-                                      ) -> Tuple[List[float], List[float]]:
+def get_training_channel_mean_stddevs(
+        hdf5_path: str,
+        data_fraction: float = 0.1,
+        seed: Optional[int] = None) -> Tuple[List[float], List[float]]:
     """
     Get the averaged channel mean and std from training data, following ImageNet
     implementation from https://github.com/soumith/imagenet-multiGPU.torch/blob/master/donkey.lua # pylint: disable=line-too-long
@@ -126,9 +126,8 @@ def get_user_num_images(hdf5_path: str, partition: str) -> Dict[str, int]:
     """
     with h5py.File(hdf5_path, 'r') as h5:
         user_num_images = {}
-        for key in tqdm.tqdm(
-                h5[f'/{partition}'].keys(),
-                desc=f"Creating user_num_images for {partition}"):
+        for key in tqdm.tqdm(h5[f'/{partition}'].keys(),
+                             desc=f"Creating user_num_images for {partition}"):
             user_num_images[key] = h5[f'/{partition}/{key}/image_ids'].shape[0]
     return user_num_images
 
@@ -165,7 +164,6 @@ def make_central_datasets(hdf5_path: str, partition: str,
     inputs_all = np.vstack(inputs_all)
     targets_all = np.vstack(targets_all)
     data_tensors = [inputs_all, targets_all]
-    return Dataset(
-        raw_data=data_tensors,
-        train_kwargs={"eval": False},
-        eval_kwargs={"eval": True})
+    return Dataset(raw_data=data_tensors,
+                   train_kwargs={"eval": False},
+                   eval_kwargs={"eval": True})
diff --git a/benchmarks/dataset/flair/download_dataset.py b/benchmarks/dataset/flair/download_dataset.py
index cb17f41..6049d66 100644
--- a/benchmarks/dataset/flair/download_dataset.py
+++ b/benchmarks/dataset/flair/download_dataset.py
@@ -30,8 +30,8 @@
 
 def extract_tar(compressed_path: str, dataset_dir: str,
                 keep_archive_after_decompress: bool):
-    subprocess.run(
-        f"tar -zxf {compressed_path} -C {dataset_dir}".split(), check=True)
+    subprocess.run(f"tar -zxf {compressed_path} -C {dataset_dir}".split(),
+                   check=True)
     if not keep_archive_after_decompress:
         os.remove(compressed_path)
 
@@ -50,26 +50,22 @@ def decompress_images(dataset_dir: str, keep_archive_after_decompress: bool):
 
 
 if __name__ == '__main__':
-    logging.basicConfig(
-        stream=sys.stdout,
-        level=logging.INFO,
-        format='%(asctime)s %(levelname)s: %(message)s')
+    logging.basicConfig(stream=sys.stdout,
+                        level=logging.INFO,
+                        format='%(asctime)s %(levelname)s: %(message)s')
 
     parser = argparse.ArgumentParser(
         description='Download the images and labels of FLAIR dataset.')
-    parser.add_argument(
-        "--dataset_dir",
-        required=True,
-        help="Path to directory of dataset to be downloaded")
-    parser.add_argument(
-        "--download_raw",
-        action="store_true",
-        help="Whether to download the raw images, "
-        "which need storage space ~1.2TB")
-    parser.add_argument(
-        "--keep_archive_after_decompress",
-        action="store_true",
-        help="Whether to keep the image tarball archives")
+    parser.add_argument("--dataset_dir",
+                        required=True,
+                        help="Path to directory of dataset to be downloaded")
+    parser.add_argument("--download_raw",
+                        action="store_true",
+                        help="Whether to download the raw images, "
+                        "which need storage space ~1.2TB")
+    parser.add_argument("--keep_archive_after_decompress",
+                        action="store_true",
+                        help="Whether to keep the image tarball archives")
     arguments = parser.parse_args()
     os.makedirs(arguments.dataset_dir, exist_ok=True)
 
diff --git a/benchmarks/dataset/flair/numpy.py b/benchmarks/dataset/flair/numpy.py
index 418ff47..2874369 100644
--- a/benchmarks/dataset/flair/numpy.py
+++ b/benchmarks/dataset/flair/numpy.py
@@ -54,13 +54,13 @@ def make_dataset_fn(user_id):
             data_order = np.random.permutation(len(inputs))
             inputs = numpy_to_tensor(inputs[data_order])
             targets = numpy_to_tensor(targets[data_order])
-        return Dataset(
-            raw_data=[inputs, targets],
-            train_kwargs={"eval": False},
-            eval_kwargs={"eval": True})
+        return Dataset(raw_data=[inputs, targets],
+                       train_kwargs={"eval": False},
+                       eval_kwargs={"eval": True})
 
-    return FederatedDataset(
-        make_dataset_fn, sampler, user_id_to_weight=user_num_images)
+    return FederatedDataset(make_dataset_fn,
+                            sampler,
+                            user_id_to_weight=user_num_images)
 
 
 def make_artificial_federated_dataset(
@@ -110,10 +110,9 @@ def make_dataset_fn(indices):
                 targets.append(np.expand_dims(target, axis=0))
         inputs = numpy_to_tensor(np.vstack(inputs))
         targets = numpy_to_tensor(np.vstack(targets))
-        return Dataset(
-            raw_data=[inputs, targets],
-            train_kwargs={"eval": False},
-            eval_kwargs={"eval": True})
+        return Dataset(raw_data=[inputs, targets],
+                       train_kwargs={"eval": False},
+                       eval_kwargs={"eval": True})
 
     data_sampler = get_data_sampler('random', len(user_image_id))
     return ArtificialFederatedDataset(make_dataset_fn, data_sampler,
diff --git a/benchmarks/dataset/flair/prepare_dataset.py b/benchmarks/dataset/flair/prepare_dataset.py
index 97144b2..436c833 100644
--- a/benchmarks/dataset/flair/prepare_dataset.py
+++ b/benchmarks/dataset/flair/prepare_dataset.py
@@ -88,8 +88,8 @@ def preprocess_federated_dataset(image_dir: str, labels_file: str,
     fine_grained_label_counter = Counter()
     with h5py.File(output_file, 'w') as h5file:
         # Iterate through users of each partition.
-        for i, user_id in tqdm.tqdm(
-                enumerate(user_metadata), total=len(user_metadata)):
+        for i, user_id in tqdm.tqdm(enumerate(user_metadata),
+                                    total=len(user_metadata)):
             # This snippet was used to generate flair_federated_small.h5
             #if i > len(user_metadata)*0.01:
             #    break
@@ -142,8 +142,8 @@ def preprocess_federated_dataset(image_dir: str, labels_file: str,
             h5file[f'/{partition}/{user_id}/image_ids'] = np.asarray(
                 image_id_array, dtype='S')
             # Tensor with dimensions [num_images,width,height,channels]
-            h5file.create_dataset(
-                f'/{partition}/{user_id}/images', data=np.stack(image_array))
+            h5file.create_dataset(f'/{partition}/{user_id}/images',
+                                  data=np.stack(image_array))
 
             if (i + 1) % LOG_INTERVAL == 0:
                 logger.info("Processed {0}/{1} users".format(
@@ -177,8 +177,8 @@ def preprocess_central_dataset(image_dir: str, labels_file: str,
                 image_id = metadata["image_id"]
                 image = Image.open(os.path.join(image_dir, f"{image_id}.jpg"))
                 partition = metadata["partition"]
-                h5file.create_dataset(
-                    f'/{partition}/{image_id}/image', data=np.asarray(image))
+                h5file.create_dataset(f'/{partition}/{image_id}/image',
+                                      data=np.asarray(image))
                 # Encode labels as a single string, separated by delimiter |
                 h5file[
                     f'/{partition}/{image_id}/labels'] = LABEL_DELIMITER.join(
@@ -204,10 +204,9 @@ def preprocess_central_dataset(image_dir: str, labels_file: str,
 
 
 if __name__ == '__main__':
-    logging.basicConfig(
-        stream=sys.stdout,
-        level=logging.INFO,
-        format='%(asctime)s %(levelname)s: %(message)s')
+    logging.basicConfig(stream=sys.stdout,
+                        level=logging.INFO,
+                        format='%(asctime)s %(levelname)s: %(message)s')
 
     argument_parser = argparse.ArgumentParser(
         description=
@@ -222,13 +221,12 @@ def preprocess_central_dataset(image_dir: str, labels_file: str,
         required=True,
         help='Path to output HDF5 file that will be constructed by this script'
     )
-    argument_parser.add_argument(
-        '--not_group_data_by_user',
-        action='store_true',
-        default=False,
-        help='If true, do not group data by user IDs.'
-        'If false, group data by user IDs to '
-        'make suitable for federated learning.')
+    argument_parser.add_argument('--not_group_data_by_user',
+                                 action='store_true',
+                                 default=False,
+                                 help='If true, do not group data by user IDs.'
+                                 'If false, group data by user IDs to '
+                                 'make suitable for federated learning.')
     arguments = argument_parser.parse_args()
 
     image_dir = os.path.join(arguments.dataset_dir, "small_images")
diff --git a/benchmarks/dataset/flair/pytorch.py b/benchmarks/dataset/flair/pytorch.py
index 9b8df7a..01e9536 100644
--- a/benchmarks/dataset/flair/pytorch.py
+++ b/benchmarks/dataset/flair/pytorch.py
@@ -75,5 +75,6 @@ def make_federated_dataset(hdf5_path: str, partition: str,
     sampler = get_user_sampler('random', user_ids)
     flair_dataset = FLAIRDataset(hdf5_path, user_ids, partition, num_classes,
                                  use_fine_grained_labels, max_num_user_images)
-    return PyTorchFederatedDataset(
-        flair_dataset, sampler, user_id_to_weight=user_num_images)
+    return PyTorchFederatedDataset(flair_dataset,
+                                   sampler,
+                                   user_id_to_weight=user_num_images)
diff --git a/benchmarks/dataset/reddit.py b/benchmarks/dataset/reddit.py
index 4abdd59..fb6e649 100644
--- a/benchmarks/dataset/reddit.py
+++ b/benchmarks/dataset/reddit.py
@@ -68,8 +68,9 @@ def get_fraction_of_users(hdf5_path: str, partition: str,
 
         tokens_cumsum = np.cumsum(list(user_num_tokens.values()))
 
-    cutoff_index = np.searchsorted(
-        tokens_cumsum, tokens_cumsum[-1] * data_fraction, side='left')
+    cutoff_index = np.searchsorted(tokens_cumsum,
+                                   tokens_cumsum[-1] * data_fraction,
+                                   side='left')
     selected_users = users[:cutoff_index]
     print(
         f'Using {int(tokens_cumsum[cutoff_index-1])}/{int(tokens_cumsum[-1])}'
@@ -121,12 +122,13 @@ def make_dataset_fn(user_id):
 
             if include_mask:
                 # Only TF1 needs a mask tensor and manual batch padding.
-                masks = (inputs != h5['metadata/pad_symbol']).astype(
-                    np.float32)
+                masks = (inputs
+                         != h5['metadata/pad_symbol']).astype(np.float32)
 
-                inputs, targets, masks = _pad_batch(
-                    inputs, targets, masks, local_batch_size, max_sequence_len,
-                    h5['metadata/pad_symbol'])
+                inputs, targets, masks = _pad_batch(inputs, targets, masks,
+                                                    local_batch_size,
+                                                    max_sequence_len,
+                                                    h5['metadata/pad_symbol'])
                 X_lengths = np.sum(masks, axis=1)
                 data_tensors = [inputs, targets, X_lengths, masks]
             else:
@@ -143,8 +145,9 @@ def make_dataset_fn(user_id):
         data_size = int(user_num_tokens[user_name] // max_sequence_len)
         for start_index in range(0, data_size, max_data_size):
             end_index = min(data_size + 1, start_index + max_data_size)
-            user_id = UserId(
-                name=user_name, start_index=start_index, end_index=end_index)
+            user_id = UserId(name=user_name,
+                             start_index=start_index,
+                             end_index=end_index)
             if (user_id.end_index -
                     user_id.start_index) >= minimum_num_datapoints_per_user:
                 user_ids.append(user_id)
@@ -154,8 +157,9 @@ def make_dataset_fn(user_id):
         u: u.end_index - u.start_index
         for i, u in enumerate(user_ids)
     }
-    return FederatedDataset(
-        make_dataset_fn, sampler, user_id_to_weight=user_id_to_weight)
+    return FederatedDataset(make_dataset_fn,
+                            sampler,
+                            user_id_to_weight=user_id_to_weight)
 
 
 def make_central_dataset(hdf5_path: str,
@@ -206,8 +210,8 @@ def make_central_dataset(hdf5_path: str,
             max_sequence_len, h5['metadata/pad_symbol'])
 
         if include_mask:
-            X_all_lengths = np.sum(
-                inputs_all != h5['metadata/pad_symbol'], axis=1)
+            X_all_lengths = np.sum(inputs_all != h5['metadata/pad_symbol'],
+                                   axis=1)
             data_tensors = [inputs_all, targets_all, X_all_lengths, masks_all]
         else:
             data_tensors = [inputs_all, targets_all]
@@ -215,13 +219,13 @@ def make_central_dataset(hdf5_path: str,
 
 
 def make_reddit_datasets(
-        data_path: str,
-        include_mask: bool,
-        local_batch_size: int,
-        max_user_tokens: int = 1600,
-        data_fraction: float = 1.0,
-        central_data_fraction: float = 0.01,
-        minimum_num_datapoints_per_user: int = 1
+    data_path: str,
+    include_mask: bool,
+    local_batch_size: int,
+    max_user_tokens: int = 1600,
+    data_fraction: float = 1.0,
+    central_data_fraction: float = 0.01,
+    minimum_num_datapoints_per_user: int = 1
 ) -> Tuple[FederatedDataset, FederatedDataset, Dataset, Dict[str, Any]]:
     """
     Create a train and val ``ArtificialFederatedDataset`` as well as a
@@ -250,9 +254,10 @@ def make_reddit_datasets(
     # Federated evaluation with `val_cohort_size>=200` is reliable enough.
     # This central evaluation is solely to compare with the same validation set
     # as a centrally trained model.
-    central_data = make_central_dataset(
-        data_path, 'val', include_mask, central_data_fraction,
-        local_batch_size, metadata['max_sequence_length'])
+    central_data = make_central_dataset(data_path, 'val', include_mask,
+                                        central_data_fraction,
+                                        local_batch_size,
+                                        metadata['max_sequence_length'])
 
     return (training_federated_dataset, val_federated_dataset, central_data,
             metadata)
diff --git a/benchmarks/dataset/stackoverflow/download_preprocess.py b/benchmarks/dataset/stackoverflow/download_preprocess.py
index 66254a8..0409c62 100644
--- a/benchmarks/dataset/stackoverflow/download_preprocess.py
+++ b/benchmarks/dataset/stackoverflow/download_preprocess.py
@@ -56,8 +56,9 @@ def _process_vocabulary(vocabulary_size, h5):
     h5['/metadata/pad_symbol'] = vocabulary[PAD]
     for token, id_ in vocabulary.items():
         # Wrap inside " because of the '.' token.
-        h5.create_dataset(
-            f'/metadata/vocabulary/"{token}"', data=id_, dtype='int32')
+        h5.create_dataset(f'/metadata/vocabulary/"{token}"',
+                          data=id_,
+                          dtype='int32')
     return vocabulary
 
 
@@ -73,6 +74,7 @@ def tokens_to_word_ids(tokens, vocab):
 
 def _make_worker_fn(tff_dataset, partition, vocabulary, max_sequence_length,
                     h5_path, lock):
+
     def _process_user(user_id):
 
         # tf Dataset with sentences from user.
@@ -95,14 +97,12 @@ def _process_user(user_id):
         with lock:
             with h5py.File(h5_path, 'a') as h5:
                 # Store encoded inputs.
-                h5.create_dataset(
-                    f'/{partition}/{user_id}/inputs',
-                    data=padded_sentences_ids[:, :-1])
+                h5.create_dataset(f'/{partition}/{user_id}/inputs',
+                                  data=padded_sentences_ids[:, :-1])
 
                 # Store encoded targets.
-                h5.create_dataset(
-                    f'/{partition}/{user_id}/targets',
-                    data=padded_sentences_ids[:, 1:])
+                h5.create_dataset(f'/{partition}/{user_id}/targets',
+                                  data=padded_sentences_ids[:, 1:])
         return user_id, token_count
 
     def _worker_fn(work_queue, result_queue):
@@ -170,15 +170,15 @@ def dl_preprocess_and_dump_h5(output_dir: str, vocabulary_size: int,
         work_queue = mp.Queue()
         results_queue = mp.Queue()
         processes = [
-            mp.Process(
-                target=_make_worker_fn(tff_dataset, partition, vocabulary,
-                                       max_sequence_length, h5_path, lock),
-                args=(work_queue, results_queue)) for _ in range(num_processes)
+            mp.Process(target=_make_worker_fn(tff_dataset, partition,
+                                              vocabulary, max_sequence_length,
+                                              h5_path, lock),
+                       args=(work_queue, results_queue))
+            for _ in range(num_processes)
         ]
         processes.append(
-            mp.Process(
-                target=_collect_num_tokens,
-                args=(results_queue, user_num_tokens, num_processes)))
+            mp.Process(target=_collect_num_tokens,
+                       args=(results_queue, user_num_tokens, num_processes)))
 
         for p in processes:
             p.start()
diff --git a/benchmarks/dataset/stackoverflow/numpy.py b/benchmarks/dataset/stackoverflow/numpy.py
index ec7924d..7f05bb6 100644
--- a/benchmarks/dataset/stackoverflow/numpy.py
+++ b/benchmarks/dataset/stackoverflow/numpy.py
@@ -52,8 +52,9 @@ def get_fraction_of_users(hdf5_path: str, partition: str,
 
         tokens_cumsum = np.cumsum(list(user_num_tokens.values()))
 
-    cutoff_index = np.searchsorted(
-        tokens_cumsum, tokens_cumsum[-1] * data_fraction, side='left')
+    cutoff_index = np.searchsorted(tokens_cumsum,
+                                   tokens_cumsum[-1] * data_fraction,
+                                   side='left')
     selected_users = users[:cutoff_index]
     print(
         f'Using {int(tokens_cumsum[cutoff_index-1])}/{int(tokens_cumsum[-1])}'
@@ -73,8 +74,8 @@ def make_tensors_fn(hdf5_path, partition, max_user_sentences, user_id):
     """
     with h5py.File(hdf5_path, 'r') as h5:
         inputs = np.array(h5[f'{partition}/{user_id}/inputs'], dtype=np.int32)
-        targets = np.array(
-            h5[f'{partition}/{user_id}/targets'], dtype=np.int32)
+        targets = np.array(h5[f'{partition}/{user_id}/targets'],
+                           dtype=np.int32)
         data_order = np.random.permutation(len(inputs))
         inputs = inputs[data_order][:max_user_sentences]
         targets = targets[data_order][:max_user_sentences]
@@ -84,10 +85,9 @@ def make_tensors_fn(hdf5_path, partition, max_user_sentences, user_id):
 def make_dataset_fn(hdf5_path, partition, max_user_sentences, user_id):
     tensors = make_tensors_fn(hdf5_path, partition, max_user_sentences,
                               user_id)
-    return Dataset(
-        raw_data=tensors,
-        train_kwargs={"eval": False},
-        eval_kwargs={"eval": True})
+    return Dataset(raw_data=tensors,
+                   train_kwargs={"eval": False},
+                   eval_kwargs={"eval": True})
 
 
 def make_federated_dataset(hdf5_path: str,
@@ -120,8 +120,9 @@ def make_federated_dataset(hdf5_path: str,
     make_dataset_fn_ = partial(make_dataset_fn, hdf5_path, partition,
                                max_user_sentences)
     sampler = get_user_sampler('random', user_ids)
-    return FederatedDataset(
-        make_dataset_fn_, sampler, user_id_to_weight=user_id_to_weight)
+    return FederatedDataset(make_dataset_fn_,
+                            sampler,
+                            user_id_to_weight=user_id_to_weight)
 
 
 def make_central_dataset(hdf5_path: str, partition: str,
@@ -156,15 +157,15 @@ def make_central_dataset(hdf5_path: str, partition: str,
             targets_all.extend(targets.tolist())
         inputs_all, targets_all = (np.array(inputs_all, dtype=np.int32),
                                    np.array(targets_all, dtype=np.int32))
-    return Dataset(
-        raw_data=[inputs_all, targets_all], eval_kwargs={"eval": True})
+    return Dataset(raw_data=[inputs_all, targets_all],
+                   eval_kwargs={"eval": True})
 
 
 def make_stackoverflow_datasets(
-        data_path: str,
-        max_user_sentences: int = 1000,
-        data_fraction: float = 1.0,
-        central_data_fraction: float = 0.01,
+    data_path: str,
+    max_user_sentences: int = 1000,
+    data_fraction: float = 1.0,
+    central_data_fraction: float = 0.01,
 ) -> Tuple[FederatedDataset, FederatedDataset, Dataset, Dict[str, Any]]:
     """
     Create a train and test ``FederatedDataset`` as well as a
diff --git a/benchmarks/dataset/stackoverflow/pytorch.py b/benchmarks/dataset/stackoverflow/pytorch.py
index bcd395b..34cb095 100644
--- a/benchmarks/dataset/stackoverflow/pytorch.py
+++ b/benchmarks/dataset/stackoverflow/pytorch.py
@@ -67,14 +67,13 @@ def make_federated_dataset(
     dataset = _PTDatasetWrap(hdf5_path, partition, max_user_sentences,
                              user_ids)
     sampler = get_user_sampler('random', user_ids)
-    return PyTorchFederatedDataset(
-        dataset,
-        sampler,
-        user_id_to_weight=user_id_to_weight,
-        num_workers=4,
-        pin_memory=True,
-        prefetch_factor=4,
-        persistent_workers=False)
+    return PyTorchFederatedDataset(dataset,
+                                   sampler,
+                                   user_id_to_weight=user_id_to_weight,
+                                   num_workers=4,
+                                   pin_memory=True,
+                                   prefetch_factor=4,
+                                   persistent_workers=False)
 
 
 def make_central_dataset(hdf5_path: str, partition: str,
@@ -106,12 +105,12 @@ def make_central_dataset(hdf5_path: str, partition: str,
 
 
 def make_stackoverflow_datasets(
-        data_path: str,
-        max_user_sentences: int = 1000,
-        data_fraction: float = 1.0,
-        central_data_fraction: float = 0.01,
-) -> Tuple[PyTorchFederatedDataset, PyTorchFederatedDataset, Dataset,
-           Dict[str, Any]]:
+    data_path: str,
+    max_user_sentences: int = 1000,
+    data_fraction: float = 1.0,
+    central_data_fraction: float = 0.01,
+) -> Tuple[PyTorchFederatedDataset, PyTorchFederatedDataset, Dataset, Dict[
+        str, Any]]:
     """
     Create a train and test ``PyTorchFederatedDataset`` as well as a
     central dataset for StackOverflow dataset.
diff --git a/benchmarks/dataset/stackoverflow/tensorflow_v2.py b/benchmarks/dataset/stackoverflow/tensorflow_v2.py
index 97a670a..5035f0f 100644
--- a/benchmarks/dataset/stackoverflow/tensorflow_v2.py
+++ b/benchmarks/dataset/stackoverflow/tensorflow_v2.py
@@ -56,11 +56,10 @@ def make_tf_dataset_fn(data):
         data = data.prefetch(10)
         return data
 
-    return TFFederatedDataset(
-        make_tf_dataset_fn,
-        sampler,
-        user_id_dtype=tf.string,
-        user_id_to_weight=user_id_to_weight)
+    return TFFederatedDataset(make_tf_dataset_fn,
+                              sampler,
+                              user_id_dtype=tf.string,
+                              user_id_to_weight=user_id_to_weight)
 
 
 def make_central_dataset(hdf5_path: str, partition: str,
@@ -92,10 +91,10 @@ def make_central_dataset(hdf5_path: str, partition: str,
 
 
 def make_stackoverflow_datasets(
-        data_path: str,
-        max_user_sentences: int = 1000,
-        data_fraction: float = 1.0,
-        central_data_fraction: float = 0.01,
+    data_path: str,
+    max_user_sentences: int = 1000,
+    data_fraction: float = 1.0,
+    central_data_fraction: float = 0.01,
 ) -> Tuple[TFFederatedDataset, TFFederatedDataset, Dataset, Dict[str, Any]]:
     """
     Create a train and test ``TFFederatedDataset`` as well as a
diff --git a/benchmarks/flair/argument_parsing.py b/benchmarks/flair/argument_parsing.py
index 69087a8..5430f78 100644
--- a/benchmarks/flair/argument_parsing.py
+++ b/benchmarks/flair/argument_parsing.py
@@ -4,11 +4,10 @@
 
 
 def add_flair_training_arguments(argument_parser):
-    argument_parser.add_argument(
-        '--central_optimizer',
-        choices=['adam', 'sgd'],
-        default='adam',
-        help='Central optimizer.')
+    argument_parser.add_argument('--central_optimizer',
+                                 choices=['adam', 'sgd'],
+                                 default='adam',
+                                 help='Central optimizer.')
 
     argument_parser.add_argument(
         '--weight_decay',
diff --git a/benchmarks/flair/train.py b/benchmarks/flair/train.py
index 4f6bfd2..04375bd 100644
--- a/benchmarks/flair/train.py
+++ b/benchmarks/flair/train.py
@@ -16,10 +16,12 @@
 
 from dataset.argument_parsing import add_dataset_arguments, get_datasets
 from model.argument_parsing import add_model_arguments, get_model_pytorch
-from utils.argument_parsing import (
-    add_algorithm_arguments, add_filepath_arguments, add_seed_arguments,
-    add_dnn_training_arguments, add_mechanism_arguments, get_algorithm,
-    parse_mechanism, maybe_inject_arguments_from_config)
+from utils.argument_parsing import (add_algorithm_arguments,
+                                    add_filepath_arguments, add_seed_arguments,
+                                    add_dnn_training_arguments,
+                                    add_mechanism_arguments, get_algorithm,
+                                    parse_mechanism,
+                                    maybe_inject_arguments_from_config)
 from utils.callback.pytorch import CentralLRDecay
 from utils.logging import init_logging
 from .argument_parsing import add_flair_training_arguments
@@ -104,15 +106,13 @@ def to_tensor(values: np.ndarray) -> torch.Tensor:
             arguments.learning_rate,
             weight_decay=arguments.weight_decay)
 
-    model = PyTorchModel(
-        model=pytorch_model,
-        local_optimizer_create=torch.optim.SGD,
-        central_optimizer=central_optimizer)
+    model = PyTorchModel(model=pytorch_model,
+                         local_optimizer_create=torch.optim.SGD,
+                         central_optimizer=central_optimizer)
 
-    backend = SimulatedBackend(
-        training_data=training_federated_dataset,
-        val_data=val_federated_dataset,
-        postprocessors=[local_privacy, central_privacy])
+    backend = SimulatedBackend(training_data=training_federated_dataset,
+                               val_data=val_federated_dataset,
+                               postprocessors=[local_privacy, central_privacy])
 
     algorithm, algorithm_params = get_algorithm(arguments)
 
@@ -127,19 +127,17 @@ def to_tensor(values: np.ndarray) -> torch.Tensor:
 
     # Central evaluation on dev data.
     callbacks = [
-        CentralEvaluationCallback(
-            central_data,
-            model_eval_params=model_eval_params,
-            frequency=arguments.evaluation_frequency),
+        CentralEvaluationCallback(central_data,
+                                  model_eval_params=model_eval_params,
+                                  frequency=arguments.evaluation_frequency),
         StopwatchCallback(),
         ModelCheckpointingCallback('./checkpoints'),
         AggregateMetricsToDisk('./metrics.csv'),
-        CentralLRDecay(
-            arguments.learning_rate,
-            0.02,
-            arguments.central_num_iterations,
-            30,
-            linear_warmup=True)
+        CentralLRDecay(arguments.learning_rate,
+                       0.02,
+                       arguments.central_num_iterations,
+                       30,
+                       linear_warmup=True)
     ]
 
     if arguments.restore_model_path is not None:
@@ -147,13 +145,12 @@ def to_tensor(values: np.ndarray) -> torch.Tensor:
         logger.info('Restored model from {}'.format(
             arguments.restore_model_path))
 
-    model = algorithm.run(
-        algorithm_params=algorithm_params,
-        backend=backend,
-        model=model,
-        model_train_params=model_train_params,
-        model_eval_params=model_eval_params,
-        callbacks=callbacks)
+    model = algorithm.run(algorithm_params=algorithm_params,
+                          backend=backend,
+                          model=model,
+                          model_train_params=model_train_params,
+                          model_eval_params=model_eval_params,
+                          callbacks=callbacks)
 
 
 if __name__ == '__main__':
diff --git a/benchmarks/image_classification/pytorch/train.py b/benchmarks/image_classification/pytorch/train.py
index 2fa8a92..6a0ce25 100644
--- a/benchmarks/image_classification/pytorch/train.py
+++ b/benchmarks/image_classification/pytorch/train.py
@@ -17,10 +17,12 @@
 from pfl.internal.ops.pytorch_ops import to_tensor
 from dataset.argument_parsing import add_dataset_arguments, get_datasets
 from model.argument_parsing import add_model_arguments, get_model_pytorch
-from utils.argument_parsing import (
-    add_algorithm_arguments, add_filepath_arguments, add_seed_arguments,
-    add_dnn_training_arguments, add_mechanism_arguments, get_algorithm,
-    parse_mechanism, maybe_inject_arguments_from_config)
+from utils.argument_parsing import (add_algorithm_arguments,
+                                    add_filepath_arguments, add_seed_arguments,
+                                    add_dnn_training_arguments,
+                                    add_mechanism_arguments, get_algorithm,
+                                    parse_mechanism,
+                                    maybe_inject_arguments_from_config)
 from utils.logging import init_logging
 
 
@@ -72,10 +74,10 @@ def main():
 
     params = [p for p in pytorch_model.parameters() if p.requires_grad]
 
-    model = PyTorchModel(
-        model=pytorch_model,
-        local_optimizer_create=torch.optim.SGD,
-        central_optimizer=torch.optim.SGD(params, arguments.learning_rate))
+    model = PyTorchModel(model=pytorch_model,
+                         local_optimizer_create=torch.optim.SGD,
+                         central_optimizer=torch.optim.SGD(
+                             params, arguments.learning_rate))
 
     weighting_strategy = WeightByDatapoints(
     ) if arguments.weight_by_samples else WeightByUser()
@@ -85,10 +87,9 @@ def main():
         CentrallyAppliedPrivacyMechanism(central_privacy)
     ]
 
-    backend = SimulatedBackend(
-        training_data=training_federated_dataset,
-        val_data=val_federated_dataset,
-        postprocessors=postprocessors)
+    backend = SimulatedBackend(training_data=training_federated_dataset,
+                               val_data=val_federated_dataset,
+                               postprocessors=postprocessors)
 
     algorithm, algorithm_params = get_algorithm(arguments)
 
@@ -112,18 +113,17 @@ def main():
         logger.info('Restored model from {}'.format(
             arguments.restore_model_path))
 
-    model = algorithm.run(
-        algorithm_params=algorithm_params,
-        backend=backend,
-        model=model,
-        model_train_params=model_train_params,
-        model_eval_params=model_eval_params,
-        callbacks=[
-            StopwatchCallback(),
-            central_evaluation_cb,
-            ModelCheckpointingCallback('./checkpoints'),
-            AggregateMetricsToDisk('./metrics.csv'),
-        ])
+    model = algorithm.run(algorithm_params=algorithm_params,
+                          backend=backend,
+                          model=model,
+                          model_train_params=model_train_params,
+                          model_eval_params=model_eval_params,
+                          callbacks=[
+                              StopwatchCallback(),
+                              central_evaluation_cb,
+                              ModelCheckpointingCallback('./checkpoints'),
+                              AggregateMetricsToDisk('./metrics.csv'),
+                          ])
 
 
 if __name__ == '__main__':
diff --git a/benchmarks/image_classification/tf/train.py b/benchmarks/image_classification/tf/train.py
index 6756dd1..a1c7de3 100644
--- a/benchmarks/image_classification/tf/train.py
+++ b/benchmarks/image_classification/tf/train.py
@@ -18,10 +18,12 @@
 
 from dataset.argument_parsing import add_dataset_arguments, get_datasets
 from model.argument_parsing import add_model_arguments, get_model_tf2
-from utils.argument_parsing import (
-    add_algorithm_arguments, add_filepath_arguments, add_seed_arguments,
-    add_dnn_training_arguments, add_mechanism_arguments, get_algorithm,
-    parse_mechanism, maybe_inject_arguments_from_config)
+from utils.argument_parsing import (add_algorithm_arguments,
+                                    add_filepath_arguments, add_seed_arguments,
+                                    add_dnn_training_arguments,
+                                    add_mechanism_arguments, get_algorithm,
+                                    parse_mechanism,
+                                    maybe_inject_arguments_from_config)
 from utils.logging import init_logging
 
 
@@ -82,10 +84,9 @@ def main():
     }
 
     central_optimizer = tf.keras.optimizers.SGD(arguments.learning_rate)
-    model = TFModel(
-        model=keras_model,
-        central_optimizer=central_optimizer,
-        metrics=metrics)
+    model = TFModel(model=keras_model,
+                    central_optimizer=central_optimizer,
+                    metrics=metrics)
 
     weighting_strategy = WeightByDatapoints(
     ) if arguments.weight_by_samples else WeightByUser()
@@ -95,10 +96,9 @@ def main():
         CentrallyAppliedPrivacyMechanism(central_privacy)
     ]
 
-    backend = SimulatedBackend(
-        training_data=training_federated_dataset,
-        val_data=val_federated_dataset,
-        postprocessors=postprocessors)
+    backend = SimulatedBackend(training_data=training_federated_dataset,
+                               val_data=val_federated_dataset,
+                               postprocessors=postprocessors)
 
     algorithm, algorithm_params = get_algorithm(arguments)
 
@@ -122,18 +122,17 @@ def main():
         logger.info('Restored model from {}'.format(
             arguments.restore_model_path))
 
-    model = algorithm.run(
-        algorithm_params=algorithm_params,
-        backend=backend,
-        model=model,
-        model_train_params=model_train_params,
-        model_eval_params=model_eval_params,
-        callbacks=[
-            StopwatchCallback(),
-            central_evaluation_cb,
-            ModelCheckpointingCallback('./checkpoints'),
-            AggregateMetricsToDisk('./metrics.csv'),
-        ])
+    model = algorithm.run(algorithm_params=algorithm_params,
+                          backend=backend,
+                          model=model,
+                          model_train_params=model_train_params,
+                          model_eval_params=model_eval_params,
+                          callbacks=[
+                              StopwatchCallback(),
+                              central_evaluation_cb,
+                              ModelCheckpointingCallback('./checkpoints'),
+                              AggregateMetricsToDisk('./metrics.csv'),
+                          ])
 
 
 if __name__ == '__main__':
diff --git a/benchmarks/lm/argument_parsing.py b/benchmarks/lm/argument_parsing.py
index c893b0c..560a7a3 100644
--- a/benchmarks/lm/argument_parsing.py
+++ b/benchmarks/lm/argument_parsing.py
@@ -12,11 +12,10 @@ def add_lm_arguments(argument_parser):
     argument_parser = add_dnn_training_arguments(argument_parser)
     argument_parser = add_mechanism_arguments(argument_parser)
 
-    argument_parser.add_argument(
-        '--central_optimizer',
-        choices=['sgd', 'adam'],
-        default='sgd',
-        help='Optimizer for central updates')
+    argument_parser.add_argument('--central_optimizer',
+                                 choices=['sgd', 'adam'],
+                                 default='sgd',
+                                 help='Optimizer for central updates')
 
     known_args, _ = argument_parser.parse_known_args()
     if known_args.central_optimizer == 'adam':
@@ -26,11 +25,10 @@ def add_lm_arguments(argument_parser):
             default=0.01,
             help='Degree of adaptivity (eps) in adaptive server optimizer.')
 
-    argument_parser.add_argument(
-        '--local_lr_decay',
-        action=store_bool,
-        default=False,
-        help='Use linear local learning rate decay')
+    argument_parser.add_argument('--local_lr_decay',
+                                 action=store_bool,
+                                 default=False,
+                                 help='Use linear local learning rate decay')
 
     argument_parser.add_argument(
         '--central_lr_num_warmup_iterations',
diff --git a/benchmarks/lm/pytorch/train.py b/benchmarks/lm/pytorch/train.py
index 2c861bc..cb85d18 100644
--- a/benchmarks/lm/pytorch/train.py
+++ b/benchmarks/lm/pytorch/train.py
@@ -15,10 +15,11 @@
 from pfl.privacy import CentrallyAppliedPrivacyMechanism
 from dataset.argument_parsing import add_dataset_arguments, get_datasets
 from model.argument_parsing import add_model_arguments, get_model_pytorch
-from utils.argument_parsing import (
-    add_algorithm_arguments, add_filepath_arguments, add_seed_arguments,
-    add_weighting_arguments, get_algorithm, parse_weighting_strategy,
-    parse_mechanism, maybe_inject_arguments_from_config)
+from utils.argument_parsing import (add_algorithm_arguments,
+                                    add_filepath_arguments, add_seed_arguments,
+                                    add_weighting_arguments, get_algorithm,
+                                    parse_weighting_strategy, parse_mechanism,
+                                    maybe_inject_arguments_from_config)
 from utils.callback.pytorch import CentralLRDecay
 from utils.logging import init_logging
 from ..argument_parsing import add_lm_arguments
@@ -79,19 +80,17 @@ def main():
     params = [p for p in pytorch_model.parameters() if p.requires_grad]
     if arguments.central_optimizer == 'adam':
         # Hyperparameters for stability, see S. Reddi et al. 2020 Appendix C.1.
-        central_optimizer = torch.optim.Adam(
-            params,
-            arguments.learning_rate,
-            eps=arguments.adaptivity_degree,
-            betas=(0.9, 0.99))
+        central_optimizer = torch.optim.Adam(params,
+                                             arguments.learning_rate,
+                                             eps=arguments.adaptivity_degree,
+                                             betas=(0.9, 0.99))
     else:
         assert arguments.central_optimizer == 'sgd'
         central_optimizer = torch.optim.SGD(params, arguments.learning_rate)
 
-    model = PyTorchModel(
-        model=pytorch_model,
-        local_optimizer_create=torch.optim.SGD,
-        central_optimizer=central_optimizer)
+    model = PyTorchModel(model=pytorch_model,
+                         local_optimizer_create=torch.optim.SGD,
+                         central_optimizer=central_optimizer)
 
     weighting_strategy = parse_weighting_strategy(arguments.weighting,
                                                   arguments.weight_clip)
@@ -101,10 +100,9 @@ def main():
         CentrallyAppliedPrivacyMechanism(central_privacy)
     ]
 
-    backend = SimulatedBackend(
-        training_data=training_federated_dataset,
-        val_data=val_federated_dataset,
-        postprocessors=postprocessors)
+    backend = SimulatedBackend(training_data=training_federated_dataset,
+                               val_data=val_federated_dataset,
+                               postprocessors=postprocessors)
 
     algorithm, algorithm_params = get_algorithm(arguments)
 
@@ -118,10 +116,9 @@ def main():
 
     callbacks = [
         StopwatchCallback(),
-        CentralEvaluationCallback(
-            central_data,
-            model_eval_params=model_eval_params,
-            frequency=arguments.evaluation_frequency),
+        CentralEvaluationCallback(central_data,
+                                  model_eval_params=model_eval_params,
+                                  frequency=arguments.evaluation_frequency),
         ModelCheckpointingCallback('./checkpoints'),
         AggregateMetricsToDisk('./metrics.csv'),
     ]
@@ -143,13 +140,12 @@ def main():
         logger.info('Restored model from {}'.format(
             arguments.restore_model_path))
 
-    model = algorithm.run(
-        algorithm_params=algorithm_params,
-        backend=backend,
-        model=model,
-        model_train_params=model_train_params,
-        model_eval_params=model_eval_params,
-        callbacks=callbacks)
+    model = algorithm.run(algorithm_params=algorithm_params,
+                          backend=backend,
+                          model=model,
+                          model_train_params=model_train_params,
+                          model_eval_params=model_eval_params,
+                          callbacks=callbacks)
 
 
 if __name__ == '__main__':
diff --git a/benchmarks/lm/tf/train.py b/benchmarks/lm/tf/train.py
index 8b2dd08..9f4574d 100644
--- a/benchmarks/lm/tf/train.py
+++ b/benchmarks/lm/tf/train.py
@@ -21,10 +21,11 @@
 from model.argument_parsing import add_model_arguments, get_model_tf2
 from model.tf2.metrics import (MaskedCategoricalAccuracy,
                                MaskedCategoricalCrossentropy, Perplexity)
-from utils.argument_parsing import (
-    add_algorithm_arguments, add_filepath_arguments, add_seed_arguments,
-    add_weighting_arguments, get_algorithm, parse_weighting_strategy,
-    parse_mechanism, maybe_inject_arguments_from_config)
+from utils.argument_parsing import (add_algorithm_arguments,
+                                    add_filepath_arguments, add_seed_arguments,
+                                    add_weighting_arguments, get_algorithm,
+                                    parse_weighting_strategy, parse_mechanism,
+                                    maybe_inject_arguments_from_config)
 from utils.callback.tensorflow import LocalLRDecay, CentralLRDecay
 from utils.logging import init_logging
 from ..argument_parsing import add_lm_arguments
@@ -95,8 +96,8 @@ def main():
 
     metrics = {
         'loss':
-        MaskedCategoricalCrossentropy(
-            from_logits=True, masked_tokens=[PAD, UNK]),
+        MaskedCategoricalCrossentropy(from_logits=True,
+                                      masked_tokens=[PAD, UNK]),
         'next word accuracy | including unk':
         MaskedCategoricalAccuracy(masked_tokens=[PAD]),
         'next word accuracy':
@@ -105,10 +106,9 @@ def main():
         Perplexity(from_logits=True, masked_tokens=[PAD, UNK])
     }
 
-    model = TFModel(
-        model=keras_model,
-        central_optimizer=central_optimizer,
-        metrics=metrics)
+    model = TFModel(model=keras_model,
+                    central_optimizer=central_optimizer,
+                    metrics=metrics)
 
     weighting_strategy = parse_weighting_strategy(arguments.weighting,
                                                   arguments.weight_clip)
@@ -118,10 +118,9 @@ def main():
         CentrallyAppliedPrivacyMechanism(central_privacy)
     ]
 
-    backend = SimulatedBackend(
-        training_data=training_federated_dataset,
-        val_data=val_federated_dataset,
-        postprocessors=postprocessors)
+    backend = SimulatedBackend(training_data=training_federated_dataset,
+                               val_data=val_federated_dataset,
+                               postprocessors=postprocessors)
 
     algorithm, algorithm_params = get_algorithm(arguments)
 
@@ -135,10 +134,9 @@ def main():
 
     callbacks = [
         StopwatchCallback(),
-        CentralEvaluationCallback(
-            central_data,
-            model_eval_params=model_eval_params,
-            frequency=arguments.evaluation_frequency),
+        CentralEvaluationCallback(central_data,
+                                  model_eval_params=model_eval_params,
+                                  frequency=arguments.evaluation_frequency),
         ModelCheckpointingCallback('./checkpoints'),
         AggregateMetricsToDisk('./metrics.csv'),
     ]
@@ -175,13 +173,12 @@ def main():
                 write_graph=False,
                 tensorboard_port=int(tb_port) if tb_port else None))
 
-    model = algorithm.run(
-        algorithm_params=algorithm_params,
-        backend=backend,
-        model=model,
-        model_train_params=model_train_params,
-        model_eval_params=model_eval_params,
-        callbacks=callbacks)
+    model = algorithm.run(algorithm_params=algorithm_params,
+                          backend=backend,
+                          model=model,
+                          model_train_params=model_train_params,
+                          model_eval_params=model_eval_params,
+                          callbacks=callbacks)
 
 
 if __name__ == '__main__':
diff --git a/benchmarks/model/argument_parsing.py b/benchmarks/model/argument_parsing.py
index 630b5b9..eb59dc0 100644
--- a/benchmarks/model/argument_parsing.py
+++ b/benchmarks/model/argument_parsing.py
@@ -29,31 +29,27 @@ def add_model_arguments(
     known_args, _ = parser.parse_known_args()
 
     if known_args.model_name in {'lm_lstm', 'lm_transformer'}:
-        parser.add_argument(
-            "--embedding_size",
-            type=int,
-            required=True,
-            help='Number of dimensions in embedding layer.')
+        parser.add_argument("--embedding_size",
+                            type=int,
+                            required=True,
+                            help='Number of dimensions in embedding layer.')
 
     if known_args.model_name == 'lm_lstm':
-        parser.add_argument(
-            "--num_cell_states",
-            type=int,
-            required=True,
-            help='Number of cell states in each LSTM.')
+        parser.add_argument("--num_cell_states",
+                            type=int,
+                            required=True,
+                            help='Number of cell states in each LSTM.')
 
-        parser.add_argument(
-            "--num_lstm_layers",
-            type=int,
-            required=True,
-            help='Number of stacked LSTM layers.')
+        parser.add_argument("--num_lstm_layers",
+                            type=int,
+                            required=True,
+                            help='Number of stacked LSTM layers.')
 
     if known_args.model_name == 'lm_transformer':
-        parser.add_argument(
-            "--hidden_size",
-            type=int,
-            required=True,
-            help='Number of hidde states in each Transformer.')
+        parser.add_argument("--hidden_size",
+                            type=int,
+                            required=True,
+                            help='Number of hidde states in each Transformer.')
 
         parser.add_argument(
             "--num_heads",
@@ -67,11 +63,10 @@ def add_model_arguments(
             required=True,
             help='Number of feed forward hidden states in each Transformer.')
 
-        parser.add_argument(
-            "--num_transformer_layers",
-            type=int,
-            required=True,
-            help='Number of stacked Transformer layers.')
+        parser.add_argument("--num_transformer_layers",
+                            type=int,
+                            required=True,
+                            help='Number of stacked Transformer layers.')
 
         parser.add_argument(
             '--dropout_rate',
@@ -92,16 +87,14 @@ def add_model_arguments(
             'shufflenet_v2_x1_5', 'shufflenet_v2_x2_0'
         ]
 
-        parser.add_argument(
-            '--model_type',
-            choices=_torchvision_architectures,
-            help='Model architecture.')
+        parser.add_argument('--model_type',
+                            choices=_torchvision_architectures,
+                            help='Model architecture.')
 
-        parser.add_argument(
-            '--pretrained',
-            action=store_bool,
-            default=False,
-            help='Use ImageNet pretrained vision model')
+        parser.add_argument('--pretrained',
+                            action=store_bool,
+                            default=False,
+                            help='Use ImageNet pretrained vision model')
 
     return parser
 
@@ -153,10 +146,11 @@ def get_model_tf2(args: argparse.Namespace):
         model = tf2.lm_lstm(args.embedding_size, args.num_cell_states,
                             args.num_lstm_layers, args.vocab_size)
     elif model_name == 'lm_transformer':
-        model = tf2.lm_transformer(
-            args.embedding_size, args.hidden_size, args.num_heads,
-            args.feedforward_size, args.num_transformer_layers,
-            args.vocab_size, args.max_sequence_length, args.dropout_rate)
+        model = tf2.lm_transformer(args.embedding_size, args.hidden_size,
+                                   args.num_heads, args.feedforward_size,
+                                   args.num_transformer_layers,
+                                   args.vocab_size, args.max_sequence_length,
+                                   args.dropout_rate)
     else:
         raise TypeError(f'Model {model_name} not implemented for TF2.')
 
diff --git a/benchmarks/model/numpy/metrics.py b/benchmarks/model/numpy/metrics.py
index 9b35c8a..793b43a 100644
--- a/benchmarks/model/numpy/metrics.py
+++ b/benchmarks/model/numpy/metrics.py
@@ -116,10 +116,9 @@ def from_predictions(cls, y_true: np.ndarray, y_pred: np.ndarray,
         bucket_indices = np.maximum(bucket_indices, 0).astype(int)
 
         def gather_bucket(label, bucket_index):
-            return unsorted_segment_sum(
-                data=label,
-                segment_ids=bucket_index,
-                num_segments=num_thresholds)
+            return unsorted_segment_sum(data=label,
+                                        segment_ids=bucket_index,
+                                        num_segments=num_thresholds)
 
         if multi_label:
             true_labels, false_labels = true_labels.T, false_labels.T
@@ -184,8 +183,8 @@ def false_negative(self):
 
     @property
     def true_positive_rate(self):
-        return np.nan_to_num(
-            self._true_positive / (self._true_positive + self._false_negative))
+        return np.nan_to_num(self._true_positive /
+                             (self._true_positive + self._false_negative))
 
     @property
     def false_positive_rate(self):
@@ -194,8 +193,8 @@ def false_positive_rate(self):
 
     @property
     def precision(self):
-        return np.nan_to_num(
-            self._true_positive / (self._true_positive + self._false_positive))
+        return np.nan_to_num(self._true_positive /
+                             (self._true_positive + self._false_positive))
 
     @property
     def recall(self):
@@ -264,7 +263,8 @@ def __init__(
             y_pred: Optional[np.ndarray] = None,
             multi_label: bool = False,
             num_thresholds: int = 200,
-            summation_method: AUCSummationMethod = AUCSummationMethod.INTERPOLATION,  # pylint: disable=line-too-long
+            summation_method: AUCSummationMethod = AUCSummationMethod.
+        INTERPOLATION,  # pylint: disable=line-too-long
     ):
         if num_thresholds <= 1:
             raise ValueError(
@@ -309,11 +309,10 @@ def __eq__(self, other):
 
     def __add__(self, other):
         assert isinstance(other, AUC)
-        return self.__class__(
-            self._confusion_matrix + other._confusion_matrix,
-            num_thresholds=self._num_thresholds,
-            summation_method=self._summation_method,
-            multi_label=self._multi_label)
+        return self.__class__(self._confusion_matrix + other._confusion_matrix,
+                              num_thresholds=self._num_thresholds,
+                              summation_method=self._summation_method,
+                              multi_label=self._multi_label)
 
     def _riemann_sum(self, x, y):
         # https://en.wikipedia.org/wiki/Riemann_sum
@@ -341,11 +340,10 @@ def to_vector(self) -> np.ndarray:
         return self._confusion_matrix.to_vector()
 
     def from_vector(self, vector: np.ndarray) -> 'AUC':
-        return self.__class__(
-            self._confusion_matrix.from_vector(vector),
-            num_thresholds=self._num_thresholds,
-            summation_method=self._summation_method,
-            multi_label=self._multi_label)
+        return self.__class__(self._confusion_matrix.from_vector(vector),
+                              num_thresholds=self._num_thresholds,
+                              summation_method=self._summation_method,
+                              multi_label=self._multi_label)
 
 
 class ROCAUC(AUC):
@@ -385,9 +383,8 @@ def overall_value(self):
 
             safe_p_ratio = np.where(
                 (p[:self._num_thresholds - 1] > 0) & (p[1:] > 0),
-                np.nan_to_num(
-                    p[:self._num_thresholds - 1] / np.maximum(p[1:], 0)),
-                np.ones_like(p[1:]))
+                np.nan_to_num(p[:self._num_thresholds - 1] /
+                              np.maximum(p[1:], 0)), np.ones_like(p[1:]))
 
             pr_auc_increment = np.nan_to_num(
                 prec_slope * (dtp + intercept * np.log(safe_p_ratio)) /
diff --git a/benchmarks/model/pytorch/cnn.py b/benchmarks/model/pytorch/cnn.py
index f4b78d0..2a02061 100644
--- a/benchmarks/model/pytorch/cnn.py
+++ b/benchmarks/model/pytorch/cnn.py
@@ -16,11 +16,11 @@
 
 
 def multi_label_cnn(
-        model_type: str,
-        num_outputs: int,
-        channel_mean: List[float],
-        channel_stddevs: List[float],
-        pretrained: bool,
+    model_type: str,
+    num_outputs: int,
+    channel_mean: List[float],
+    channel_stddevs: List[float],
+    pretrained: bool,
 ):
     """
     A CNN used for multi-label classification task.
@@ -55,12 +55,12 @@ class MultiLabelCNN(nn.Module):
         """
 
         def __init__(
-                self,
-                torchvision_model_type: str,
-                num_outputs: int,
-                channel_mean: List[float],
-                channel_stddevs: List[float],
-                pretrained: bool,
+            self,
+            torchvision_model_type: str,
+            num_outputs: int,
+            channel_mean: List[float],
+            channel_stddevs: List[float],
+            pretrained: bool,
         ):
             super().__init__()
             self._num_outputs = num_outputs
@@ -163,10 +163,9 @@ def metrics(self, inputs, targets, eval=True):
                 "micro recall":
                 Weighted(tps_sum, tps_sum + fns_sum),
                 "micro AP":
-                AveragedPrecision(
-                    y_true=targets.cpu().numpy(),
-                    y_pred=scores.cpu().numpy(),
-                    multi_label=False),
+                AveragedPrecision(y_true=targets.cpu().numpy(),
+                                  y_pred=scores.cpu().numpy(),
+                                  multi_label=False),
                 # Macro metrics: averaged over classes
                 "macro loss":
                 MacroWeighted(summed_loss.cpu().numpy(), num_predictions),
@@ -179,10 +178,9 @@ def metrics(self, inputs, targets, eval=True):
                 MacroWeighted(tps.cpu().numpy(),
                               tps.cpu().numpy() + fns.cpu().numpy()),
                 "macro AP":
-                AveragedPrecision(
-                    y_true=targets.cpu().numpy(),
-                    y_pred=scores.cpu().numpy(),
-                    multi_label=True),
+                AveragedPrecision(y_true=targets.cpu().numpy(),
+                                  y_pred=scores.cpu().numpy(),
+                                  multi_label=True),
             }
 
     return MultiLabelCNN(model_type, num_outputs, channel_mean,
diff --git a/benchmarks/model/pytorch/dnn.py b/benchmarks/model/pytorch/dnn.py
index 63a54c1..bbe4e0a 100644
--- a/benchmarks/model/pytorch/dnn.py
+++ b/benchmarks/model/pytorch/dnn.py
@@ -48,6 +48,6 @@ def simple_dnn(input_shape: Tuple[int, ...], num_outputs: int) -> nn.Module:
     McMahan et al. 2017 https://arxiv.org/pdf/1602.05629.pdf.
     See ``dnn`` for description about parameters.
     """
-    return functools.partial(
-        dnn, hidden_dims=[200, 200])(
-            input_shape, num_outputs=num_outputs)
+    return functools.partial(dnn, hidden_dims=[200,
+                                               200])(input_shape,
+                                                     num_outputs=num_outputs)
diff --git a/benchmarks/model/pytorch/layer.py b/benchmarks/model/pytorch/layer.py
index 2619789..27916ad 100644
--- a/benchmarks/model/pytorch/layer.py
+++ b/benchmarks/model/pytorch/layer.py
@@ -29,8 +29,8 @@ def forward(self, input: torch.Tensor) -> torch.Tensor:
         else:
             exponential_average_factor = self.momentum
 
-        bn_training = (self.running_mean is None) and (self.running_var is
-                                                       None)
+        bn_training = (self.running_mean is None) and (self.running_var
+                                                       is None)
 
         return F.batch_norm(
             input,
@@ -47,6 +47,7 @@ def forward(self, input: torch.Tensor) -> torch.Tensor:
 
 
 class FrozenBatchNorm1D(_FrozenBatchNorm):
+
     def _check_input_dim(self, input):
         if input.dim() != 2 and input.dim() != 3:
             raise ValueError('expected 2D or 3D input (got {}D input)'.format(
@@ -54,6 +55,7 @@ def _check_input_dim(self, input):
 
 
 class FrozenBatchNorm2D(_FrozenBatchNorm):
+
     def _check_input_dim(self, input):
         if input.dim() != 4:
             raise ValueError('expected 4D input (got {}D input)'.format(
@@ -61,6 +63,7 @@ def _check_input_dim(self, input):
 
 
 class FrozenBatchNorm3D(_FrozenBatchNorm):
+
     def _check_input_dim(self, input):
         if input.dim() != 5:
             raise ValueError('expected 5D input (got {}D input)'.format(
@@ -77,6 +80,7 @@ def forward(self, inputs: torch.Tensor):
 
 
 class PositionalEncoding(nn.Module):
+
     def __init__(self, embedding_size: int, max_sequence_length: int):
         super().__init__()
         pe = positional_encoding(max_sequence_length, embedding_size)
diff --git a/benchmarks/model/pytorch/lstm.py b/benchmarks/model/pytorch/lstm.py
index 04e5a2b..5496c1f 100644
--- a/benchmarks/model/pytorch/lstm.py
+++ b/benchmarks/model/pytorch/lstm.py
@@ -11,6 +11,7 @@
 
 
 class LMBase(nn.Module):
+
     def __init__(self, pad_symbol: int, unk_symbol: int):
         super().__init__()
         self._pad_symbol = pad_symbol
@@ -21,10 +22,10 @@ def loss(self, inputs, targets, eval=False):
         inputs, targets = inputs.long(), targets.long()
         logits = self.forward(inputs)
         log_probs = F.log_softmax(logits, dim=-1)
-        loss_fct = torch.nn.NLLLoss(
-            reduction='sum', ignore_index=self._pad_symbol)
-        loss = loss_fct(
-            log_probs.reshape(-1, log_probs.shape[-1]), targets.reshape(-1))
+        loss_fct = torch.nn.NLLLoss(reduction='sum',
+                                    ignore_index=self._pad_symbol)
+        loss = loss_fct(log_probs.reshape(-1, log_probs.shape[-1]),
+                        targets.reshape(-1))
         num_tokens = torch.sum(targets != self._pad_symbol)
         return loss / num_tokens
 
@@ -34,8 +35,8 @@ def metrics(self, inputs, targets, eval=False):
         inputs, targets = inputs.long(), targets.long()
         logits = self.forward(inputs)
         log_probs = F.log_softmax(logits, dim=-1)
-        loss_fct = torch.nn.NLLLoss(
-            reduction='none', ignore_index=self._pad_symbol)
+        loss_fct = torch.nn.NLLLoss(reduction='none',
+                                    ignore_index=self._pad_symbol)
 
         targets = targets.reshape(-1)
         loss = loss_fct(log_probs.reshape(-1, log_probs.shape[-1]), targets)
@@ -71,12 +72,14 @@ def metrics(self, inputs, targets, eval=False):
 
 
 class LMLSTM(LMBase):
+
     def __init__(self, embedding_size: int, num_cell_states: int,
                  num_lstm_layers: int, vocab_size: int, pad_symbol: int,
                  unk_symbol: int):
         super().__init__(pad_symbol, unk_symbol)
-        self._embeddings = nn.Embedding(
-            vocab_size, embedding_size, padding_idx=pad_symbol)
+        self._embeddings = nn.Embedding(vocab_size,
+                                        embedding_size,
+                                        padding_idx=pad_symbol)
         self._lstms = nn.ModuleList()
 
         input_size = embedding_size
diff --git a/benchmarks/model/pytorch/module_modification.py b/benchmarks/model/pytorch/module_modification.py
index 92ee3fa..7c35e62 100644
--- a/benchmarks/model/pytorch/module_modification.py
+++ b/benchmarks/model/pytorch/module_modification.py
@@ -25,9 +25,9 @@ def _replace_child(root: nn.Module, child_name: str,
 
 
 def replace_all_modules(
-        root: nn.Module,
-        target_class: Type[nn.Module],
-        converter: Callable[[nn.Module], nn.Module],
+    root: nn.Module,
+    target_class: Type[nn.Module],
+    converter: Callable[[nn.Module], nn.Module],
 ) -> nn.Module:
     """
     Converts all the submodules (of root) that have the same
@@ -55,8 +55,9 @@ def _batchnorm_to_groupnorm(
         paper *Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour*
         https://arxiv.org/pdf/1706.02677.pdf
     """
-    return nn.GroupNorm(
-        min(32, module.num_features), module.num_features, affine=True)
+    return nn.GroupNorm(min(32, module.num_features),
+                        module.num_features,
+                        affine=True)
 
 
 def _batchnorm_to_freeze_batchnorm(
@@ -76,12 +77,11 @@ def match_dim():
         elif isinstance(module, nn.BatchNorm3d):
             return FrozenBatchNorm3D
 
-    return match_dim()(
-        num_features=module.num_features,
-        eps=module.eps,
-        momentum=module.momentum,
-        affine=module.affine,
-        track_running_stats=module.track_running_stats)
+    return match_dim()(num_features=module.num_features,
+                       eps=module.eps,
+                       momentum=module.momentum,
+                       affine=module.affine,
+                       track_running_stats=module.track_running_stats)
 
 
 def validate_no_batchnorm(module: nn.Module):
@@ -96,9 +96,9 @@ def validate_no_batchnorm(module: nn.Module):
 
 
 def convert_batchnorm_modules(
-        model: nn.Module,
-        converter: Callable[[nn.modules.batchnorm._BatchNorm], nn.
-                            Module] = _batchnorm_to_groupnorm,
+    model: nn.Module,
+    converter: Callable[[nn.modules.batchnorm._BatchNorm],
+                        nn.Module] = _batchnorm_to_groupnorm,
 ) -> nn.Module:
     """
     Converts all BatchNorm modules to another module
diff --git a/benchmarks/model/pytorch/transformer.py b/benchmarks/model/pytorch/transformer.py
index 8507535..6a6fd52 100644
--- a/benchmarks/model/pytorch/transformer.py
+++ b/benchmarks/model/pytorch/transformer.py
@@ -10,22 +10,23 @@
 
 
 class LMTransformer(LMBase):
+
     def __init__(self, embedding_size: int, hidden_size: int, num_heads: int,
                  feedforward_size: int, num_transformer_layers: int,
                  dropout_rate: float, vocab_size: int,
                  max_sequence_length: int, pad_symbol: int, unk_symbol: int):
         super().__init__(pad_symbol, unk_symbol)
         self._embedding_size = embedding_size
-        self._embeddings = nn.Embedding(
-            vocab_size, embedding_size, padding_idx=pad_symbol)
+        self._embeddings = nn.Embedding(vocab_size,
+                                        embedding_size,
+                                        padding_idx=pad_symbol)
         self._positional_encoder = PositionalEncoding(embedding_size,
                                                       max_sequence_length)
-        encoder_layers = nn.TransformerEncoderLayer(
-            hidden_size,
-            num_heads,
-            feedforward_size,
-            dropout_rate,
-            batch_first=True)
+        encoder_layers = nn.TransformerEncoderLayer(hidden_size,
+                                                    num_heads,
+                                                    feedforward_size,
+                                                    dropout_rate,
+                                                    batch_first=True)
         self._transformer_encoder = nn.TransformerEncoder(
             encoder_layers, num_transformer_layers)
 
diff --git a/benchmarks/model/tf2/cnn.py b/benchmarks/model/tf2/cnn.py
index 500f6e3..c6c1a14 100644
--- a/benchmarks/model/tf2/cnn.py
+++ b/benchmarks/model/tf2/cnn.py
@@ -24,9 +24,10 @@ def simple_cnn(input_shape: Tuple[int, ...],
     """
 
     return tf.keras.models.Sequential([
-        Conv2D(
-            32, kernel_size=(3, 3), activation='relu',
-            input_shape=input_shape),
+        Conv2D(32,
+               kernel_size=(3, 3),
+               activation='relu',
+               input_shape=input_shape),
         Conv2D(64, kernel_size=(3, 3), activation='relu'),
         MaxPool2D(pool_size=(2, 2)),
         Dropout(0.25),
@@ -55,8 +56,11 @@ def _conv_norm(x, filters, kernel_size, weight_decay=1e-4, strides=(1, 1)):
     return layer
 
 
-def _conv_norm_relu(x, filters, kernel_size, weight_decay=1e-4, strides=(1,
-                                                                         1)):
+def _conv_norm_relu(x,
+                    filters,
+                    kernel_size,
+                    weight_decay=1e-4,
+                    strides=(1, 1)):
     layer = _conv_norm(x, filters, kernel_size, weight_decay, strides)
     layer = tf.nn.relu(layer)
     return layer
diff --git a/benchmarks/model/tf2/dnn.py b/benchmarks/model/tf2/dnn.py
index 15d1ab3..6d69350 100644
--- a/benchmarks/model/tf2/dnn.py
+++ b/benchmarks/model/tf2/dnn.py
@@ -26,8 +26,8 @@ def dnn(input_shape: Tuple[int, ...], hidden_dims: Tuple[int, ...],
     """
 
     return tf.keras.models.Sequential([
-        tf.keras.layers.Reshape(
-            input_shape=input_shape, target_shape=(np.prod(input_shape), )),
+        tf.keras.layers.Reshape(input_shape=input_shape,
+                                target_shape=(np.prod(input_shape), )),
         *[
             tf.keras.layers.Dense(dim, activation=tf.nn.relu)
             for dim in hidden_dims
@@ -44,6 +44,6 @@ def simple_dnn(input_shape: Tuple[int, ...],
     McMahan et al. 2017 https://arxiv.org/pdf/1602.05629.pdf.
     See ``dnn`` for description about parameters.
     """
-    return functools.partial(
-        dnn, hidden_dims=[200, 200])(
-            input_shape, num_outputs=num_outputs)
+    return functools.partial(dnn, hidden_dims=[200,
+                                               200])(input_shape,
+                                                     num_outputs=num_outputs)
diff --git a/benchmarks/model/tf2/layer.py b/benchmarks/model/tf2/layer.py
index 72c1f4e..aae71f5 100644
--- a/benchmarks/model/tf2/layer.py
+++ b/benchmarks/model/tf2/layer.py
@@ -17,11 +17,12 @@ def __init__(self, vocab_size: int, embedding_size: int,
                  max_sequence_length: int):
         super().__init__()
         self.embedding_size = embedding_size
-        self.embedding = tf.keras.layers.Embedding(
-            vocab_size, embedding_size, mask_zero=True)
+        self.embedding = tf.keras.layers.Embedding(vocab_size,
+                                                   embedding_size,
+                                                   mask_zero=True)
         self.pos_encoding = tf.convert_to_tensor(
-            positional_encoding(
-                length=max_sequence_length, depth=embedding_size))
+            positional_encoding(length=max_sequence_length,
+                                depth=embedding_size))
 
     def compute_mask(self, *args, **kwargs):
         return self.embedding.compute_mask(*args, **kwargs)
diff --git a/benchmarks/model/tf2/lstm.py b/benchmarks/model/tf2/lstm.py
index 72a9922..2c00615 100644
--- a/benchmarks/model/tf2/lstm.py
+++ b/benchmarks/model/tf2/lstm.py
@@ -27,8 +27,9 @@ def lm_lstm(embedding_size: int, num_cell_states: int, num_lstm_layers: int,
 
     inputs = tf.keras.layers.Input(shape=(None, ))
     # Embedding layer automatically masks padding for consecutive layers.
-    input_embedding = tf.keras.layers.Embedding(
-        input_dim=vocab_size, output_dim=embedding_size, mask_zero=True)
+    input_embedding = tf.keras.layers.Embedding(input_dim=vocab_size,
+                                                output_dim=embedding_size,
+                                                mask_zero=True)
     intermediate = input_embedding(inputs)
 
     for _ in range(num_lstm_layers):
@@ -38,8 +39,9 @@ def lm_lstm(embedding_size: int, num_cell_states: int, num_lstm_layers: int,
     if embedding_size != num_cell_states:
         intermediate = tf.keras.layers.Dense(embedding_size)(intermediate)
 
-    logits = tf.matmul(
-        intermediate, input_embedding.embeddings, transpose_b=True)
+    logits = tf.matmul(intermediate,
+                       input_embedding.embeddings,
+                       transpose_b=True)
 
     keras_model = tf.keras.Model(inputs=inputs, outputs=logits)
     print(keras_model.summary())
diff --git a/benchmarks/model/tf2/transformer.py b/benchmarks/model/tf2/transformer.py
index 487a618..3d8fa88 100644
--- a/benchmarks/model/tf2/transformer.py
+++ b/benchmarks/model/tf2/transformer.py
@@ -53,18 +53,18 @@ def lm_transformer(embedding_size: int,
     intermediate = tf.keras.layers.Dropout(dropout_rate)(intermediate)
 
     for _ in range(num_transformer_layers):
-        encoder_layer = CausalEncoderLayer(
-            hidden_size=hidden_size,
-            num_heads=num_heads,
-            feedforward_size=feedforward_size,
-            dropout_rate=dropout_rate)
+        encoder_layer = CausalEncoderLayer(hidden_size=hidden_size,
+                                           num_heads=num_heads,
+                                           feedforward_size=feedforward_size,
+                                           dropout_rate=dropout_rate)
         intermediate = encoder_layer(intermediate)
     # Dense layer changes dimension from rnn_layer_size to input_embedding_size
     if embedding_size != hidden_size:
         intermediate = tf.keras.layers.Dense(embedding_size)(intermediate)
 
-    logits = tf.matmul(
-        intermediate, input_embedding.embedding.embeddings, transpose_b=True)
+    logits = tf.matmul(intermediate,
+                       input_embedding.embedding.embeddings,
+                       transpose_b=True)
 
     keras_model = tf.keras.Model(inputs=inputs, outputs=logits)
     print(keras_model.summary())
diff --git a/benchmarks/poetry.lock b/benchmarks/poetry.lock
index 2e941da..77e0195 100644
--- a/benchmarks/poetry.lock
+++ b/benchmarks/poetry.lock
@@ -1,4 +1,4 @@
-# This file is automatically @generated by Poetry 1.6.1 and should not be changed by hand.
+# This file is automatically @generated by Poetry 1.7.1 and should not be changed by hand.
 
 [[package]]
 name = "absl-py"
@@ -581,6 +581,25 @@ files = [
     {file = "idna-3.4.tar.gz", hash = "sha256:814f528e8dead7d329833b91c5faa87d60bf71824cd12a7530b5526063d02cb4"},
 ]
 
+[[package]]
+name = "importlib-metadata"
+version = "7.0.1"
+description = "Read metadata from Python packages"
+optional = false
+python-versions = ">=3.8"
+files = [
+    {file = "importlib_metadata-7.0.1-py3-none-any.whl", hash = "sha256:4805911c3a4ec7c3966410053e9ec6a1fecd629117df5adee56dfc9432a1081e"},
+    {file = "importlib_metadata-7.0.1.tar.gz", hash = "sha256:f238736bb06590ae52ac1fab06a3a9ef1d8dce2b7a35b5ab329371d6c8f5d2cc"},
+]
+
+[package.dependencies]
+zipp = ">=0.5"
+
+[package.extras]
+docs = ["furo", "jaraco.packaging (>=9.3)", "jaraco.tidelift (>=1.4)", "rst.linker (>=1.9)", "sphinx (<7.2.5)", "sphinx (>=3.5)", "sphinx-lint"]
+perf = ["ipython"]
+testing = ["flufl.flake8", "importlib-resources (>=1.3)", "packaging", "pyfakefs", "pytest (>=6)", "pytest-black (>=0.3.7)", "pytest-checkdocs (>=2.4)", "pytest-cov", "pytest-enabler (>=2.2)", "pytest-mypy (>=0.9.1)", "pytest-perf (>=0.9.2)", "pytest-ruff"]
+
 [[package]]
 name = "iniconfig"
 version = "2.0.0"
@@ -1434,6 +1453,11 @@ files = [
     {file = "scikit_learn-1.3.1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f66eddfda9d45dd6cadcd706b65669ce1df84b8549875691b1f403730bdef217"},
     {file = "scikit_learn-1.3.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c6448c37741145b241eeac617028ba6ec2119e1339b1385c9720dae31367f2be"},
     {file = "scikit_learn-1.3.1-cp311-cp311-win_amd64.whl", hash = "sha256:c413c2c850241998168bbb3bd1bb59ff03b1195a53864f0b80ab092071af6028"},
+    {file = "scikit_learn-1.3.1-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:ef540e09873e31569bc8b02c8a9f745ee04d8e1263255a15c9969f6f5caa627f"},
+    {file = "scikit_learn-1.3.1-cp312-cp312-macosx_12_0_arm64.whl", hash = "sha256:9147a3a4df4d401e618713880be023e36109c85d8569b3bf5377e6cd3fecdeac"},
+    {file = "scikit_learn-1.3.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:d2cd3634695ad192bf71645702b3df498bd1e246fc2d529effdb45a06ab028b4"},
+    {file = "scikit_learn-1.3.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:0c275a06c5190c5ce00af0acbb61c06374087949f643ef32d355ece12c4db043"},
+    {file = "scikit_learn-1.3.1-cp312-cp312-win_amd64.whl", hash = "sha256:0e1aa8f206d0de814b81b41d60c1ce31f7f2c7354597af38fae46d9c47c45122"},
     {file = "scikit_learn-1.3.1-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:52b77cc08bd555969ec5150788ed50276f5ef83abb72e6f469c5b91a0009bbca"},
     {file = "scikit_learn-1.3.1-cp38-cp38-macosx_12_0_arm64.whl", hash = "sha256:a683394bc3f80b7c312c27f9b14ebea7766b1f0a34faf1a2e9158d80e860ec26"},
     {file = "scikit_learn-1.3.1-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a15d964d9eb181c79c190d3dbc2fff7338786bf017e9039571418a1d53dab236"},
@@ -2183,15 +2207,35 @@ files = [
 
 [[package]]
 name = "yapf"
-version = "0.26.0"
-description = "A formatter for Python code."
+version = "0.40.2"
+description = "A formatter for Python code"
 optional = false
-python-versions = "*"
+python-versions = ">=3.7"
 files = [
-    {file = "yapf-0.26.0-py2.py3-none-any.whl", hash = "sha256:f58069d5e0df60c078f3f986b8a63acfda73aa6ebb7cd423f6eabd1cb06420ba"},
-    {file = "yapf-0.26.0.tar.gz", hash = "sha256:edb47be90a56ca6f3075fe24f119a22225fbd62c66777b5d3916a7e9e793891b"},
+    {file = "yapf-0.40.2-py3-none-any.whl", hash = "sha256:adc8b5dd02c0143108878c499284205adb258aad6db6634e5b869e7ee2bd548b"},
+    {file = "yapf-0.40.2.tar.gz", hash = "sha256:4dab8a5ed7134e26d57c1647c7483afb3f136878b579062b786c9ba16b94637b"},
 ]
 
+[package.dependencies]
+importlib-metadata = ">=6.6.0"
+platformdirs = ">=3.5.1"
+tomli = ">=2.0.1"
+
+[[package]]
+name = "zipp"
+version = "3.17.0"
+description = "Backport of pathlib-compatible object wrapper for zip files"
+optional = false
+python-versions = ">=3.8"
+files = [
+    {file = "zipp-3.17.0-py3-none-any.whl", hash = "sha256:0e923e726174922dce09c53c59ad483ff7bbb8e572e00c7f7c46b88556409f31"},
+    {file = "zipp-3.17.0.tar.gz", hash = "sha256:84e64a1c28cf7e91ed2078bb8cc8c259cb19b76942096c8d7b84947690cabaf0"},
+]
+
+[package.extras]
+docs = ["furo", "jaraco.packaging (>=9.3)", "jaraco.tidelift (>=1.4)", "rst.linker (>=1.9)", "sphinx (<7.2.5)", "sphinx (>=3.5)", "sphinx-lint"]
+testing = ["big-O", "jaraco.functools", "jaraco.itertools", "more-itertools", "pytest (>=6)", "pytest-black (>=0.3.7)", "pytest-checkdocs (>=2.4)", "pytest-cov", "pytest-enabler (>=2.2)", "pytest-ignore-flaky", "pytest-mypy (>=0.9.1)", "pytest-ruff"]
+
 [extras]
 pytorch = ["pfl", "pfl", "torch", "torch", "torchvision", "torchvision"]
 tf = ["pfl", "pfl", "tensorflow", "tensorflow_addons"]
@@ -2199,4 +2243,4 @@ tf = ["pfl", "pfl", "tensorflow", "tensorflow_addons"]
 [metadata]
 lock-version = "2.0"
 python-versions = ">=3.10,<3.11"
-content-hash = "9cb359884c1afef3134ade071648270036107d35b9d23fcc75affdce946992bd"
+content-hash = "e0e978dbb0c039df57e29160c2087de0c5ea389422a7ce5a15b9943728fe8b18"
diff --git a/benchmarks/pyproject.toml b/benchmarks/pyproject.toml
index bebed77..0bd9ec7 100644
--- a/benchmarks/pyproject.toml
+++ b/benchmarks/pyproject.toml
@@ -42,9 +42,7 @@ pytest-lazy-fixture = "^0.6.3"
 mypy = "1.5"
 ruff = "0.0.290"
 scikit-learn = "^1.0.2"
-# TODO: upgrade yapf in future PR
-#yapf = "^0.40.1"
-yapf = "^0.26.0"
+yapf = "^0.40.1"
 
 # TODO: these pytorch sources are only needed while pfl is installed from github address.
 [[tool.poetry.source]]
diff --git a/benchmarks/utils/__init__.py b/benchmarks/utils/__init__.py
index 583969f..b8cee13 100644
--- a/benchmarks/utils/__init__.py
+++ b/benchmarks/utils/__init__.py
@@ -1,4 +1,3 @@
 # -*- coding: utf-8 -*-
 
 # Copyright © 2023-2024 Apple Inc.
-
diff --git a/benchmarks/utils/argument_parsing.py b/benchmarks/utils/argument_parsing.py
index 04e2c4c..02acfdc 100644
--- a/benchmarks/utils/argument_parsing.py
+++ b/benchmarks/utils/argument_parsing.py
@@ -17,10 +17,11 @@
 from pfl.aggregate.weighting import WeightByUser
 from pfl.algorithm import (NNAlgorithmParams, FederatedAveraging, FedProx,
                            FedProxParams)
-from pfl.privacy import (
-    CentrallyApplicablePrivacyMechanism, CentrallyAppliedPrivacyMechanism,
-    NoPrivacy, GaussianMechanism, LaplaceMechanism, GaussianMechanism,
-    NormClippingOnly, PrivUnitMechanism, PLDPrivacyAccountant)
+from pfl.privacy import (CentrallyApplicablePrivacyMechanism,
+                         CentrallyAppliedPrivacyMechanism, NoPrivacy,
+                         GaussianMechanism, LaplaceMechanism,
+                         GaussianMechanism, NormClippingOnly,
+                         PrivUnitMechanism, PLDPrivacyAccountant)
 from pfl.privacy.ftrl_mechanism import BandedMatrixFactorizationMechanism
 
 logger = logging.getLogger(name=__name__)
@@ -46,8 +47,10 @@ def maybe_inject_arguments_from_config():
 def add_seed_arguments(
         parser: argparse.ArgumentParser) -> argparse.ArgumentParser:
 
-    parser.add_argument(
-        "--seed", type=int, default=0, help=('Seed to use for pseudo rng'))
+    parser.add_argument("--seed",
+                        type=int,
+                        default=0,
+                        help=('Seed to use for pseudo rng'))
 
     return parser
 
@@ -59,15 +62,13 @@ def add_filepath_arguments(
         '--args_config',
         help='Path to YAML configuration file containing the arguments.')
 
-    parser.add_argument(
-        '--data_path',
-        required=True,
-        help='The path from which the dataset will be read.')
+    parser.add_argument('--data_path',
+                        required=True,
+                        help='The path from which the dataset will be read.')
 
-    parser.add_argument(
-        "--restore_model_path",
-        default=None,
-        help='Path to model checkpoint to restore')
+    parser.add_argument("--restore_model_path",
+                        default=None,
+                        help='Path to model checkpoint to restore')
 
     return parser
 
@@ -113,6 +114,7 @@ def parse_weighting_strategy(name, weight_clip):
 
 
 class store_bool(argparse.Action):
+
     def __init__(self, option_strings, dest, **kwargs):
         argparse.Action.__init__(self, option_strings, dest, **kwargs)
 
@@ -251,6 +253,7 @@ def add_mechanism_arguments(argument_parser):
         'applied.')
 
     class OrderAction(argparse.Action):
+
         def __init__(self, option_strings, dest, nargs=None, **kwargs):
             if nargs is not None:
                 raise ValueError("nargs not allowed")
@@ -430,14 +433,13 @@ def add_algorithm_arguments(
     specified in `algorithm_name` argument.
     """
 
-    parser.add_argument(
-        '--algorithm_name',
-        choices=[
-            'fedavg',
-            'fedprox',
-        ],
-        default='fedavg',
-        help='Which algorithm to train with.')
+    parser.add_argument('--algorithm_name',
+                        choices=[
+                            'fedavg',
+                            'fedprox',
+                        ],
+                        default='fedavg',
+                        help='Which algorithm to train with.')
 
     # Get the value of `algorithm_name` argument and dynamically add
     # arguments depending on which algorithm is chosen.
diff --git a/benchmarks/utils/callback/pytorch.py b/benchmarks/utils/callback/pytorch.py
index 91f79d1..a917b68 100644
--- a/benchmarks/utils/callback/pytorch.py
+++ b/benchmarks/utils/callback/pytorch.py
@@ -21,6 +21,7 @@ def polynomial_lr_lambda(current_step: int, lr_init: float, lr_end: float,
 
 
 class CentralLRDecay(TrainingProcessCallback):
+
     def __init__(self,
                  init_learning_rate: float,
                  end_learning_rate: float,
diff --git a/benchmarks/utils/callback/tensorflow.py b/benchmarks/utils/callback/tensorflow.py
index e70ca04..5ebad37 100644
--- a/benchmarks/utils/callback/tensorflow.py
+++ b/benchmarks/utils/callback/tensorflow.py
@@ -10,6 +10,7 @@
 
 
 class LocalLRDecay(TrainingProcessCallback):
+
     def __init__(self, initial_learning_rate, central_num_iterations):
         # Linear warmup over 20 central iterations.
         self._num_warmup_iterations = 20
@@ -39,6 +40,7 @@ def after_central_iteration(self, central_iteration, config,
 
 
 class CentralLRDecay(TrainingProcessCallback):
+
     def __init__(self,
                  init_learning_rate: float,
                  end_learning_rate: float,
diff --git a/benchmarks/utils/weighting.py b/benchmarks/utils/weighting.py
index aca9636..ee4fcb0 100644
--- a/benchmarks/utils/weighting.py
+++ b/benchmarks/utils/weighting.py
@@ -7,6 +7,7 @@
 
 
 class WeightByTokens(WeightByDatapoints):
+
     def _num_tokens(self, user_dataset):
         return np.sum(user_dataset.raw_data[0] != 0)
 
@@ -15,21 +16,25 @@ def reweight_statistics(self, statistics, user_dataset):
 
 
 class WeightBySqrtTokens(WeightByTokens):
+
     def reweight_statistics(self, statistics, user_dataset):
         statistics.reweight(np.sqrt(self._num_tokens(user_dataset)))
 
 
 class WeightByCubeRootTokens(WeightByTokens):
+
     def reweight_statistics(self, statistics, user_dataset):
         statistics.reweight(self._num_tokens(user_dataset)**(1 / 3))
 
 
 class WeightByLogTokens(WeightByTokens):
+
     def reweight_statistics(self, statistics, user_dataset):
         statistics.reweight(np.log(self._num_tokens(user_dataset)))
 
 
 class WeightByTokensClipped(WeightByTokens):
+
     def __init__(self, weight_clip):
         self._weight_clip = weight_clip
 


Test patch:
diff --git a/benchmarks/test/conftest.py b/benchmarks/test/conftest.py
index fe1c340..695ed9c 100644
--- a/benchmarks/test/conftest.py
+++ b/benchmarks/test/conftest.py
@@ -5,11 +5,10 @@
 
 
 def pytest_addoption(parser):
-    parser.addoption(
-        '--macos',
-        action='store_true',
-        default=False,
-        help='run tests that require MacOS')
+    parser.addoption('--macos',
+                     action='store_true',
+                     default=False,
+                     help='run tests that require MacOS')
 
 
 def pytest_collection_modifyitems(config, items):
diff --git a/benchmarks/test/dataset/cifar10/test_download_preprocess.py b/benchmarks/test/dataset/cifar10/test_download_preprocess.py
index f1407a4..38cfe47 100644
--- a/benchmarks/test/dataset/cifar10/test_download_preprocess.py
+++ b/benchmarks/test/dataset/cifar10/test_download_preprocess.py
@@ -49,6 +49,7 @@ def mock_endpoint():
 
 
 class TestDownloadPreprocess:
+
     def test_dl_preprocess_and_dump(self, mock_endpoint, tmp_path):
         dl_preprocess_and_dump(tmp_path)
         with open(os.path.join(tmp_path, 'cifar10_train.p'), 'rb') as f:
diff --git a/benchmarks/test/dataset/stackoverflow/test_download_preprocess.py b/benchmarks/test/dataset/stackoverflow/test_download_preprocess.py
index e76c7ab..8d2f97b 100644
--- a/benchmarks/test/dataset/stackoverflow/test_download_preprocess.py
+++ b/benchmarks/test/dataset/stackoverflow/test_download_preprocess.py
@@ -17,6 +17,7 @@
 
 @pytest.fixture
 def mock_tffdatas():
+
     def make_client(user_id):
         sentence_data = {'tokens': tf.constant("a b c")}
         # Number of mock sentences == user_id value
@@ -36,6 +37,7 @@ def make_mock_tffdata(partition_user_id):
 
 @pytest.fixture(autouse=True)
 def tff(mock_tffdatas):
+
     def mock_word_counts(vocab_size):
         # vocabulary is alphabet.
         return {chr(97 + i): i for i in range(vocab_size)}
@@ -54,6 +56,7 @@ def mock_word_counts(vocab_size):
 # Only run if TF2 is installed. Using tf.data.
 @pytest.mark.skipif(get_tf_major_version() < 2, reason='not tf>=2')
 class TestDownloadPreprocess:
+
     def test_get_vocabulary(self, tff):
         from dataset.stackoverflow.download_preprocess import get_vocabulary
 
diff --git a/benchmarks/test/dataset/stackoverflow/test_stackoverflow.py b/benchmarks/test/dataset/stackoverflow/test_stackoverflow.py
index bfe48a7..52e6e66 100644
--- a/benchmarks/test/dataset/stackoverflow/test_stackoverflow.py
+++ b/benchmarks/test/dataset/stackoverflow/test_stackoverflow.py
@@ -36,10 +36,10 @@ def h5_path(tmp_path):
             for user_id in range(1, 4):
                 features = np.zeros((3, 2)) + user_id + partition_ix
                 targets = np.zeros((3, 2)) + user_id * 10 + partition_ix
-                h5.create_dataset(
-                    f'/{partition}/{user_id}/inputs', data=features)
-                h5.create_dataset(
-                    f'/{partition}/{user_id}/targets', data=targets)
+                h5.create_dataset(f'/{partition}/{user_id}/inputs',
+                                  data=features)
+                h5.create_dataset(f'/{partition}/{user_id}/targets',
+                                  data=targets)
 
     yield h5_path
 
@@ -93,6 +93,7 @@ def _check_user_dataset(dataset, ix, partition_ix, expected_num_users):
 
 
 class TestStackOverflowNumpy:
+
     def test_get_metadata(self, h5_path):
         metadata = get_metadata(h5_path)
         assert metadata == {
@@ -103,8 +104,9 @@ def test_get_metadata(self, h5_path):
         }
 
     def test_get_fraction_of_users(self, h5_path):
-        users = get_fraction_of_users(
-            h5_path, partition='train', data_fraction=0.7)
+        users = get_fraction_of_users(h5_path,
+                                      partition='train',
+                                      data_fraction=0.7)
         assert users == ['1', '2']
 
     def test_get_user_weights(self, h5_path):
@@ -122,23 +124,22 @@ def test_get_user_weights(self, h5_path):
         pytest.param(
             pytest.lazy_fixture('tf2_so_module'),
             marks=[
-                pytest.mark.skipif(
-                    get_tf_major_version() < 2, reason='not tf>=2'),
+                pytest.mark.skipif(get_tf_major_version() < 2,
+                                   reason='not tf>=2'),
                 # This test makes pytest get stuck when running in CI, only run
                 # on MacOs for now while rdar://107677363 is not fixed.
                 pytest.mark.macos
             ],
             id='tf_v2'),
-        pytest.param(
-            pytest.lazy_fixture('pytorch_so_module'),
-            marks=[
-                pytest.mark.skipif(
-                    not get_pytorch_major_version(),
-                    reason='PyTorch not installed'),
-            ],
-            id='pytorch'),
+        pytest.param(pytest.lazy_fixture('pytorch_so_module'),
+                     marks=[
+                         pytest.mark.skipif(not get_pytorch_major_version(),
+                                            reason='PyTorch not installed'),
+                     ],
+                     id='pytorch'),
     ])
 class TestStackOverflowAll:
+
     @pytest.mark.parametrize('data_fraction', [1.0, 0.7])
     @pytest.mark.parametrize('partition', ['train', 'val'])
     def test_make_federated_dataset(self, so_module, partition, data_fraction,
@@ -187,9 +188,13 @@ def test_make_stackoverflow_datasets(self, so_module, h5_path):
         }
         for ix in range(3):
             user, _ = next(fed_train)
-            _check_user_dataset(
-                user, ix=0, partition_ix=0, expected_num_users=1)
+            _check_user_dataset(user,
+                                ix=0,
+                                partition_ix=0,
+                                expected_num_users=1)
         for ix in range(3):
             user, _ = next(fed_val)
-            _check_user_dataset(
-                user, ix=ix, partition_ix=1, expected_num_users=3)
+            _check_user_dataset(user,
+                                ix=ix,
+                                partition_ix=1,
+                                expected_num_users=3)
diff --git a/benchmarks/test/model/numpy/test_metrics.py b/benchmarks/test/model/numpy/test_metrics.py
index 4c212e4..12b72d7 100644
--- a/benchmarks/test/model/numpy/test_metrics.py
+++ b/benchmarks/test/model/numpy/test_metrics.py
@@ -57,10 +57,10 @@ def _get_sklearn_auc(auc_type):
 
 def _get_keras_auc(auc_type, multi_label):
     from tensorflow.keras.metrics import AUC as KerasAUC  # type: ignore
-    keras_roc_auc = lambda multi_label: KerasAUC(
-        multi_label=multi_label, curve="roc")
-    keras_pr_auc = lambda multi_label: KerasAUC(
-        multi_label=multi_label, curve="pr")
+    keras_roc_auc = lambda multi_label: KerasAUC(multi_label=multi_label,
+                                                 curve="roc")
+    keras_pr_auc = lambda multi_label: KerasAUC(multi_label=multi_label,
+                                                curve="pr")
 
     if auc_type == "roc":
         return keras_roc_auc(multi_label)
@@ -71,6 +71,7 @@ def _get_keras_auc(auc_type, multi_label):
 
 
 class TestAUC:
+
     @pytest.mark.parametrize('auc_type', ["roc", "pr", "ap"])
     @pytest.mark.parametrize('multi_label', [True, False])
     def test_auc_aggregation(self, auc_type, multi_label):
@@ -82,8 +83,9 @@ def test_auc_aggregation(self, auc_type, multi_label):
         for _ in range(num_workers):
             y_score = _get_random_probabilities()
             y_true = np.round(_get_random_probabilities()).astype(int)
-            worker_auc = _get_bucket_auc(auc_type)(
-                y_true=y_true, y_pred=y_score, multi_label=multi_label)
+            worker_auc = _get_bucket_auc(auc_type)(y_true=y_true,
+                                                   y_pred=y_score,
+                                                   multi_label=multi_label)
             if aggregated_auc is None:
                 aggregated_auc = worker_auc
             else:
@@ -120,10 +122,9 @@ def test_auc_against_sklearn(self, auc_type, multi_label):
                     y_true=y_true,
                     y_score=y_score,
                     average="macro" if multi_label else "micro")
-                assert np.isclose(
-                    bucket_value, sklearn_value, atol=5e-3), (
-                        f"Our {auc_type} AUC value: {bucket_value},"
-                        f" sklearn {auc_type} AUC value: {sklearn_value}")
+                assert np.isclose(bucket_value, sklearn_value, atol=5e-3), (
+                    f"Our {auc_type} AUC value: {bucket_value},"
+                    f" sklearn {auc_type} AUC value: {sklearn_value}")
 
     @pytest.mark.skipif(get_tf_major_version() < 2, reason='not tf>=2')
     @pytest.mark.parametrize('auc_type', ["roc", "pr"])
diff --git a/benchmarks/test/model/pytorch/test_cnn.py b/benchmarks/test/model/pytorch/test_cnn.py
index 78d6ca7..e96733f 100644
--- a/benchmarks/test/model/pytorch/test_cnn.py
+++ b/benchmarks/test/model/pytorch/test_cnn.py
@@ -14,9 +14,10 @@
     import torch
 
 
-@pytest.mark.skipif(
-    not get_pytorch_major_version(), reason='PyTorch not installed')
+@pytest.mark.skipif(not get_pytorch_major_version(),
+                    reason='PyTorch not installed')
 class TestPyTorchCNN:
+
     @pytest.mark.parametrize('input_shape', [(32, 32, 3), (28, 28, 1)])
     @pytest.mark.parametrize('output_shape', [10, 100])
     def test_output_shape(self, input_shape: Tuple[int, ...],
diff --git a/benchmarks/test/model/pytorch/test_dnn.py b/benchmarks/test/model/pytorch/test_dnn.py
index dac9418..f3b7bbc 100644
--- a/benchmarks/test/model/pytorch/test_dnn.py
+++ b/benchmarks/test/model/pytorch/test_dnn.py
@@ -14,9 +14,10 @@
     import torch
 
 
-@pytest.mark.skipif(
-    not get_pytorch_major_version(), reason='PyTorch not installed')
+@pytest.mark.skipif(not get_pytorch_major_version(),
+                    reason='PyTorch not installed')
 class TestPyTorchDNN:
+
     @pytest.mark.parametrize('input_shape', [(32, 32, 3), (28, 28, 1)])
     @pytest.mark.parametrize('output_shape', [10, 100])
     def test_output_shape(self, input_shape: Tuple[int, ...],
diff --git a/benchmarks/test/model/pytorch/test_lstm.py b/benchmarks/test/model/pytorch/test_lstm.py
index d67c5fe..21a8508 100644
--- a/benchmarks/test/model/pytorch/test_lstm.py
+++ b/benchmarks/test/model/pytorch/test_lstm.py
@@ -14,9 +14,10 @@
 _UNK_SYMBOL = 1
 
 
-@pytest.mark.skipif(
-    not get_pytorch_major_version(), reason='PyTorch not installed')
+@pytest.mark.skipif(not get_pytorch_major_version(),
+                    reason='PyTorch not installed')
 class TestPyTorchCNN:
+
     @pytest.mark.parametrize('embedding_size', [64, 128])
     @pytest.mark.parametrize('num_cell_states', [64, 128])
     @pytest.mark.parametrize('num_lstm_layers', [1, 2])
@@ -29,8 +30,9 @@ def test_output_shape(self, embedding_size: int, num_cell_states: int,
                                 num_lstm_layers, vocab_size, _PAD_SYMBOL,
                                 _UNK_SYMBOL)
         batch_size = 3
-        inputs = np.random.randint(
-            0, vocab_size, size=(batch_size, sequence_length))
+        inputs = np.random.randint(0,
+                                   vocab_size,
+                                   size=(batch_size, sequence_length))
         logits = pytorch_model(torch.from_numpy(inputs))
         assert logits.shape == (batch_size, sequence_length, vocab_size)
 
diff --git a/benchmarks/test/test_argument_parsing.py b/benchmarks/test/test_argument_parsing.py
index 5837846..b811a3f 100644
--- a/benchmarks/test/test_argument_parsing.py
+++ b/benchmarks/test/test_argument_parsing.py
@@ -12,6 +12,7 @@
 
 
 class TestStoreBool(unittest.TestCase):
+
     def test_values(self):
         parser = argparse.ArgumentParser()
         parser.add_argument('--argument_name', action=store_bool)
@@ -38,8 +39,9 @@ def test_failure(self):
 
     def test_default(self):
         parser = argparse.ArgumentParser()
-        parser.add_argument(
-            '--argument_name', action=store_bool, default=False)
+        parser.add_argument('--argument_name',
+                            action=store_bool,
+                            default=False)
 
         result = parser.parse_args([])
         self.assertEqual(result.argument_name, False)
diff --git a/tests/data/test_federated_dataset.py b/tests/data/test_federated_dataset.py
index 31d2be1..00acfe4 100644
--- a/tests/data/test_federated_dataset.py
+++ b/tests/data/test_federated_dataset.py
@@ -422,21 +422,26 @@ def test_3_workers(self, mock_get_ops, make_fed_data):
                                               [70 + i * 30])
 
 
-@pytest.mark.parametrize('make_fed_data', [
-    pytest.param(lazy_fixture('make_fed_data_numpy'), id='numpy'),
-    pytest.param(lazy_fixture('make_fed_data_pytorch'),
-                 marks=[
-                     pytest.mark.skipif(not get_pytorch_major_version(),
-                                        reason='PyTorch not installed')
-                 ],
-                 id='pytorch'),
-    pytest.param(lazy_fixture('make_fed_data_tf'),
-                 marks=[
-                     pytest.mark.skipif(get_tf_major_version() < 2,
-                                        reason='not tf>=2'),
-                 ],
-                 id='tensorflow')
-])
+@pytest.mark.parametrize(
+    'make_fed_data',
+    [
+        pytest.param(lazy_fixture('make_fed_data_numpy'), id='numpy'),
+        pytest.param(lazy_fixture('make_fed_data_pytorch'),
+                     marks=[
+                         pytest.mark.skipif(not get_pytorch_major_version(),
+                                            reason='PyTorch not installed')
+                     ],
+                     id='pytorch'),
+        # TODO: Other tests get stuck when this is run with pytest. Need to fix.
+        # It is related to using multiprocess in the generator of
+        # tf.data.Dataset.from_generator
+        #pytest.param(lazy_fixture('make_fed_data_tf'),
+        #             marks=[
+        #                 pytest.mark.skipif(get_tf_major_version() < 2,
+        #                                    reason='not tf>=2'),
+        #             ],
+        #             id='tensorflow')
+    ])
 @pytest.mark.parametrize('user_id_to_weight', (True, ), indirect=True)
 class TestFrameworkFederatedDatasetSorted: