Repository: lishenghui/blades (language: python)

Issue:
Update README.rst



Discussion / hints:

Fix patch:
diff --git a/.gitignore b/.gitignore
index 8525839..e232b4f 100644
--- a/.gitignore
+++ b/.gitignore
@@ -17,3 +17,5 @@ src/blades/examples/data/*s
 */build/
 /docs/source/examples/
 /docs/source/_static/
+/src/dist/
+*.pytest_cache/
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
new file mode 100644
index 0000000..dcbbb8c
--- /dev/null
+++ b/.pre-commit-config.yaml
@@ -0,0 +1,50 @@
+repos:
+#  - repo: https://github.com/PyCQA/isort
+#    rev: 5.10.1
+#    hooks:
+#      - id: isort
+  - repo: https://github.com/psf/black
+    rev: 22.8.0
+    hooks:
+      - id: black
+  - repo: https://github.com/PyCQA/flake8
+    rev: 5.0.4
+    hooks:
+      - id: flake8
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v4.3.0
+    hooks:
+      - id: trailing-whitespace
+      - id: check-yaml
+      - id: end-of-file-fixer
+      - id: requirements-txt-fixer
+#      - id: double-quote-string-fixer
+      - id: check-merge-conflict
+      - id: fix-encoding-pragma
+        args: ["--remove"]
+      - id: mixed-line-ending
+        args: ["--fix=lf"]
+  - repo: https://github.com/codespell-project/codespell
+    rev: v2.2.1
+    hooks:
+      - id: codespell
+  - repo: https://github.com/executablebooks/mdformat
+    rev: 0.7.16
+    hooks:
+      - id: mdformat
+        args: ["--number"]
+        additional_dependencies:
+          - mdformat-openmmlab
+          - mdformat_frontmatter
+          - linkify-it-py
+  - repo: https://github.com/myint/docformatter
+    rev: v1.5.0
+    hooks:
+      - id: docformatter
+        args: ["--in-place", "--wrap-descriptions", "79"]
+#  - repo: https://github.com/open-mmlab/pre-commit-hooks
+#    rev: v0.2.0  # Use the ref you want to point at
+#    hooks:
+#      - id: check-algo-readme
+#      - id: check-copyright
+#        args: ["mmdet"]  # replace the dir_to_check with your expected directory to check
diff --git a/config/example.yaml b/config/example.yaml
index 0d31acb..ab4a26c 100644
--- a/config/example.yaml
+++ b/config/example.yaml
@@ -1,5 +1,5 @@
 
-dataset: mnist
+dataset: cifar100
 non_iid: False
 alpha: 0.1
 num_clients: 50
@@ -10,8 +10,8 @@ num_byzantine: 5
 attack: ALIE
 attack_kws: {"num_clients": 50, "num_byzantine": 5}
 
-agg: MultiKrum
-aggregator_kws: {num_excluded: 5, k: 5}
+agg: Mean
+#aggregator_kws: {num_excluded: 1, k: 5}
 #aggregator_kws: {"linkage": "average"}
 
 
@@ -34,6 +34,5 @@ dp_clip_threshold: 0.5
 
 # If you try to debug the code, use `local_mode:True`
 local_mode: False
-num_actors: 4
+num_actors: 20
 gpu_per_actor: 0.2
-
diff --git a/docs/source/attacks.md b/docs/source/attacks.md
index bfbc3af..27c85a7 100644
--- a/docs/source/attacks.md
+++ b/docs/source/attacks.md
@@ -1,32 +1,34 @@
-
 (sec:attacks)=
 
 # Attacks
+
 (sec:buildinattacks)=
 
-## Build-in Attacks
-The following attack strategies are currently implemented in ``Blades``:
-### Untargeted Attack
+## Built-in Attacks
 
-| Strategy          | Description                                                                                                                                           | Examples                                                                                                        |
-| ---------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------- | --------------------------------------------------------------------------------------------------------------- |
-| **Noise** | Put random noise to the updates. | [[**Example**]](https://github.com/bladesteam/blades/blob/master/src/blades/attackers/noiseclient.py) |
-| **Labelflipping** | *Fang et al.* [Local Model Poisoning Attacks to Byzantine-Robust Federated Learning](https://www.usenix.org/conference/usenixsecurity20/presentation/fang), *USENIX Security' 20* | [[**Example**]](https://github.com/bladesteam/blades/blob/master/src/blades/attackers/labelflippingclient.py) |
-| **Signflipping** | *Li et al.* [RSA: Byzantine-Robust Stochastic Aggregation Methods for Distributed Learning from Heterogeneous Datasets](https://ojs.aaai.org/index.php/AAAI/article/view/3968), *AAAI' 19* | [[**Example**]](https://github.com/bladesteam/blades/blob/master/src/blades/attackers/signflippingclient.py) |
-| **ALIE** | *Baruch et al.* [A little is enough: Circumventing defenses for distributed learning](https://proceedings.neurips.cc/paper/2019/hash/ec1c59141046cd1866bbbcdfb6ae31d4-Abstract.html), *NeurIPS' 19* | [[**Example**]](https://github.com/bladesteam/blades/blob/master/src/blades/attackers/alieclient.py) |
-| **IPM** | *Xie et al.* [Fall of empires: Breaking byzantine- tolerant sgd by inner product manipulation](https://arxiv.org/abs/1903.03936), *UAI' 20* | [[**Example**]](https://github.com/bladesteam/blades/blob/master/src/blades/attackers/ipmclient.py) |
+The following attack strategies are currently implemented in `Blades`:
 
+### Untargeted Attack
+
+| Strategy          | Description                                                                                                                                                                                         | Examples                                                                                                        |
+| ----------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | --------------------------------------------------------------------------------------------------------------- |
+| **Noise**         | Put random noise to the updates.                                                                                                                                                                    | [\[**Example**\]](https://github.com/bladesteam/blades/blob/master/src/blades/attackers/noiseclient.py)         |
+| **Labelflipping** | *Fang et al.* [Local Model Poisoning Attacks to Byzantine-Robust Federated Learning](https://www.usenix.org/conference/usenixsecurity20/presentation/fang), *USENIX Security' 20*                   | [\[**Example**\]](https://github.com/bladesteam/blades/blob/master/src/blades/attackers/labelflippingclient.py) |
+| **Signflipping**  | *Li et al.* [RSA: Byzantine-Robust Stochastic Aggregation Methods for Distributed Learning from Heterogeneous Datasets](https://ojs.aaai.org/index.php/AAAI/article/view/3968), *AAAI' 19*          | [\[**Example**\]](https://github.com/bladesteam/blades/blob/master/src/blades/attackers/signflippingclient.py)  |
+| **ALIE**          | *Baruch et al.* [A little is enough: Circumventing defenses for distributed learning](https://proceedings.neurips.cc/paper/2019/hash/ec1c59141046cd1866bbbcdfb6ae31d4-Abstract.html), *NeurIPS' 19* | [\[**Example**\]](https://github.com/bladesteam/blades/blob/master/src/blades/attackers/alieclient.py)          |
+| **IPM**           | *Xie et al.* [Fall of empires: Breaking byzantine- tolerant sgd by inner product manipulation](https://arxiv.org/abs/1903.03936), *UAI' 20*                                                         | [\[**Example**\]](https://github.com/bladesteam/blades/blob/master/src/blades/attackers/ipmclient.py)           |
 
 ## Customizing Attacks
+
 (sec:customattacks)=
 The following example shows how to customize attack strategy.
 
 ```python
 import ray
-from blades.client import ByzantineClient
+from blades.core.client import ByzantineClient
 from blades.datasets import CIFAR10
 from blades.models.cifar10 import CCTNet
-from blades.simulator import Simulator
+from blades.core.simulator import Simulator
 
 cifar10 = CIFAR10(num_clients=20, iid=True)  # built-in federated cifar10 dataset
 
@@ -34,7 +36,7 @@ cifar10 = CIFAR10(num_clients=20, iid=True)  # built-in federated cifar10 datase
 class MaliciousClient(ByzantineClient):
     def __init__(self, *args, **kwargs):
         super(ByzantineClient).__init__(*args, **kwargs)
-    
+
     def omniscient_callback(self, simulator):
         updates = []
         for w in simulator._clients:
@@ -60,7 +62,7 @@ simulator.register_attackers(attackers)
 
 # runtime parameters
 run_params = {
-    "model": CCTNet(),  # global model
+    "global_model": CCTNet(),  # global global_model
     "server_optimizer": 'SGD',  # server optimizer
     "client_optimizer": 'SGD',  # client optimizer
     "loss": "crossentropy",  # loss function
@@ -70,4 +72,4 @@ run_params = {
 }
 simulator.run(**run_params)
 
-```
\ No newline at end of file
+```
diff --git a/docs/source/conf.py b/docs/source/conf.py
index b848cea..a7c84a0 100644
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -16,19 +16,19 @@
 # import sphinx_rtd_theme
 
 # -- Project information -----------------------------------------------------
+import os
+import sys
 
-project = 'Blades'
-copyright = '2022, Blades Team'
-author = 'Blades Team'
-
-
+project = "Blades"
+copyright = "2022, Blades Team"
+author = "Blades Team"
 
 intersphinx_mapping = {
-    'python': ('https://docs.python.org/3/', None),
-    'numpy': ('https://numpy.org/doc/stable/', None),
-    'scipy': ('https://docs.scipy.org/doc/scipy/', None),
-    'numba': ('https://numba.readthedocs.io/en/stable', None),
-    'torch': ('https://pytorch.org/docs/stable', None),
+    "python": ("https://docs.python.org/3/", None),
+    "numpy": ("https://numpy.org/doc/stable/", None),
+    "scipy": ("https://docs.scipy.org/doc/scipy/", None),
+    "numba": ("https://numba.readthedocs.io/en/stable", None),
+    "torch": ("https://pytorch.org/docs/stable", None),
 }
 
 # -- General configuration ---------------------------------------------------
@@ -38,45 +38,46 @@
 # ones.
 napoleon_use_param = True
 
-extensions = ['sphinx.ext.autodoc',
-              'sphinx.ext.autosummary',
-              'sphinx.ext.doctest',
-              'sphinx.ext.extlinks',
-              'sphinx.ext.intersphinx',
-              'sphinx.ext.mathjax',
-              'sphinx.ext.napoleon',
-              'sphinx_gallery.gen_gallery',
-              # 'sphinx.ext.linkcode',
-              'sphinx_copybutton',
-              'sphinx.ext.viewcode',
-#               'sphinx_autodoc_typehints',
-#               # 'myst_parser',
-              "nbsphinx",
-              'sphinx_gallery.load_style',
-#               'sphinx_rtd_theme',
-              'm2r2',
-              # 'sphinx.ext.pngmath',
-              ]
+extensions = [
+    "sphinx.ext.autodoc",
+    "sphinx.ext.autosummary",
+    "sphinx.ext.doctest",
+    "sphinx.ext.extlinks",
+    "sphinx.ext.intersphinx",
+    "sphinx.ext.mathjax",
+    "sphinx.ext.napoleon",
+    "sphinx_gallery.gen_gallery",
+    # 'sphinx.ext.linkcode',
+    "sphinx_copybutton",
+    "sphinx.ext.viewcode",
+    #               'sphinx_autodoc_typehints',
+    #               # 'myst_parser',
+    "nbsphinx",
+    "sphinx_gallery.load_style",
+    #               'sphinx_rtd_theme',
+    "m2r2",
+    # 'sphinx.ext.pngmath',
+]
 
 # Add any paths that contain templates here, relative to this directory.
-templates_path = ['_templates']
+templates_path = ["_templates"]
 
 # List of patterns, relative to source directory, that match files and
 # directories to ignore when looking for source files.
 # This pattern also affects html_static_path and html_extra_path.
 exclude_patterns = []
 
-html_logo = '_static/blades_logo.png'
+html_logo = "_static/blades_logo.png"
 # The name of an image file (within the static path) to use as favicon of the
 # docs.  This file should be a Windows icon file (.ico) being 16x16 or 32x32
 # pixels large.
-html_favicon = '_static/favicon.ico'
+html_favicon = "_static/favicon.ico"
 
 sphinx_gallery_conf = {
-    'line_numbers': False,
-    'ignore_pattern': '/todo_',
-    'examples_dirs': '../../src/blades/examples',   # path to your example scripts
-    'gallery_dirs': 'examples',  # path to where to save gallery generated output
+    "line_numbers": False,
+    "ignore_pattern": "/todo_",
+    "examples_dirs": "../../src/blades/examples",  # path to your example scripts
+    "gallery_dirs": "examples",  # path to where to save gallery generated output
 }
 
 # -- Options for HTML output -------------------------------------------------
@@ -86,26 +87,22 @@
 #
 # html_theme = 'sphinx_rtd_theme'
 # Add any paths that contain custom static files (such as style sheets) here,
-html_theme = 'pydata_sphinx_theme'
+html_theme = "pydata_sphinx_theme"
 # relative to this directory. They are copied after the builtin static files,
 # so a file named "default.css" will override the builtin "default.css".
-html_static_path = ['_static']
-
+html_static_path = ["_static"]
 
 html_theme_options = {
-    'navigation_depth': 5,
-    'collapse_navigation': False,
+    "navigation_depth": 5,
+    "collapse_navigation": False,
     "icon_links": [
-            {
-                "name": "GitHub",
-                "url": "https://github.com/bladesteam/blades",
-                "icon": "fab fa-github-square",
-            },
-        ],
+        {
+            "name": "GitHub",
+            "url": "https://github.com/bladesteam/blades",
+            "icon": "fab fa-github-square",
+        },
+    ],
 }
 
-import os
-import sys
-
-sys.path.insert(0, os.path.abspath('/Users/sheli564/Desktop/blades/src/blades'))
-sys.path.insert(0, os.path.abspath('/Users/sheli564/Desktop/blades/'))
+sys.path.insert(0, os.path.abspath("/Users/sheli564/Desktop/blades/src/blades"))
+sys.path.insert(0, os.path.abspath("/Users/sheli564/Desktop/blades/"))
diff --git a/docs/source/defenses.md b/docs/source/defenses.md
index f1ff4df..a188d43 100644
--- a/docs/source/defenses.md
+++ b/docs/source/defenses.md
@@ -34,7 +34,7 @@ import ray
 import torch
 from blades.datasets import CIFAR10
 from blades.models.cifar10 import CCTNet
-from blades.simulator import Simulator
+from blades.core.simulator import Simulator
 
 cifar10 = CIFAR10(num_clients=20, iid=True)  # built-in federated cifar10 dataset
 
@@ -45,7 +45,7 @@ class Median():
         values_upper, _ = stacked.median(dim=0)
         values_lower, _ = (-stacked).median(dim=0)
         return (values_upper - values_lower) / 2
-    
+
 
 # configuration parameters
 conf_params = {
diff --git a/docs/source/references/aggregators.rst b/docs/source/references/aggregators.rst
index 9853f9e..25b1389 100644
--- a/docs/source/references/aggregators.rst
+++ b/docs/source/references/aggregators.rst
@@ -6,7 +6,7 @@ Aggregators
     :show-inheritance:
     :member-order: bysource
 
-.. automodule:: aggregators.krum
+.. automodule:: aggregators.multikrum
     :members:
     :show-inheritance:
     :member-order: bysource
@@ -50,4 +50,4 @@ Aggregators
 .. automodule:: aggregators.fltrust
     :members:
     :show-inheritance:
-    :member-order: bysource
\ No newline at end of file
+    :member-order: bysource
diff --git a/docs/source/references/client.rst b/docs/source/references/client.rst
index ce22f10..87abc8b 100644
--- a/docs/source/references/client.rst
+++ b/docs/source/references/client.rst
@@ -2,7 +2,7 @@
 Client
 ===========
 
-.. automodule:: client    
+.. automodule:: core.client
     :members:
     :show-inheritance:
-    :member-order: bysource
\ No newline at end of file
+    :member-order: bysource
diff --git a/docs/source/references/server.rst b/docs/source/references/server.rst
index 187fd30..2a697d3 100644
--- a/docs/source/references/server.rst
+++ b/docs/source/references/server.rst
@@ -3,7 +3,7 @@ Server
 ===========
 
 
-.. automodule:: server    
+.. automodule:: core.server
     :members:
     :show-inheritance:
-    :member-order: bysource
\ No newline at end of file
+    :member-order: bysource
diff --git a/docs/source/references/simulator.rst b/docs/source/references/simulator.rst
index d959d89..b489b3c 100644
--- a/docs/source/references/simulator.rst
+++ b/docs/source/references/simulator.rst
@@ -1,4 +1,4 @@
 Simulator
 ==========
-.. automodule:: simulator    
-    :members:
\ No newline at end of file
+.. automodule:: core.simulator
+    :members:
diff --git a/requirements.txt b/requirements.txt
index dedf92e..edff672 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,14 +1,15 @@
-torch>=1.10.2
-torchvision>=0.11.3
+blades>=0.0.13
+matplotlib>=3.4.1
 numpy>=1.19.4
-scipy>=1.5.4
+pre-commit
 ray>=1.0.0
-sklearn>=0.0
-scikit-learn>=1.0.2
-matplotlib>=3.4.1
 requests>=2.27.1
-blades>=0.0.13
+ruamel.yaml>=0.17.21
+ruamel.yaml.clib>=0.2.6
+scikit-learn>=1.0.2
+scipy>=1.5.4
 setuptools>=58.0.4
-ruamel.yaml==0.17.21
-ruamel.yaml.clib==0.2.6
-tqdm
\ No newline at end of file
+sklearn>=0.0
+torch>=1.10.2
+torchvision>=0.11.3
+tqdm
diff --git a/scripts/args.py b/scripts/args.py
index 7c321fe..4f88337 100644
--- a/scripts/args.py
+++ b/scripts/args.py
@@ -1,8 +1,11 @@
 import argparse
 import os
-from blades.utils.utils import over_write_args_from_file
+
 import torch
 
+from blades.utils.utils import over_write_args_from_file
+
+
 def parse_arguments():
     parser = argparse.ArgumentParser()
     parser.add_argument("--non_iid", action="store_true", default=False)
@@ -14,11 +17,20 @@ def parse_arguments():
     parser.add_argument("--batch_size", type=int, default=32)
     parser.add_argument("--test_batch_size", type=int, default=128)
     parser.add_argument("--validate_interval", type=int, default=100)
-    parser.add_argument("--trusted_id", type=int, default=49)
-    parser.add_argument("--attack", type=str, default='signflipping', help="Select attack types.")
-    parser.add_argument("--dataset", type=str, default='cifar10', help="Dataset")
-    parser.add_argument("--algorithm", type=str, default='fedsgd', help="Optimization algorithm, either 'fedavg' or 'fedsgd'.")
-    parser.add_argument("--agg", type=str, default='clippedclustering', help="Aggregator.")
+    parser.add_argument("--trusted_id", type=int, default=None)
+    parser.add_argument(
+        "--attack", type=str, default="signflipping", help="Select attack types."
+    )
+    parser.add_argument("--dataset", type=str, default="cifar10", help="Dataset")
+    parser.add_argument(
+        "--algorithm",
+        type=str,
+        default="fedsgd",
+        help="Optimization algorithm, either 'fedavg' or 'fedsgd'.",
+    )
+    parser.add_argument(
+        "--agg", type=str, default="clippedclustering", help="Aggregator."
+    )
     parser.add_argument("--lr", type=float, default=0.1, help="learning rate")
     parser.add_argument("--alpha", type=float, default=0.1)
     parser.add_argument("--num_clients", type=int, default=20)
@@ -26,87 +38,107 @@ def parse_arguments():
 
     parser.add_argument("--num_actors", type=int, default=5)
     parser.add_argument("--gpu_per_actor", type=float, default=0.2)
-    
+
     parser.add_argument("--dp", action="store_true", default=False)
 
-    # Parameters for DP. They will take effect only if `dp` is set to be `True`.
+    # Parameters for DP. They will take effect only if `dp`
+    # is  `True`.
     parser.add_argument("--dp_privacy_delta", type=float, default=1e-6)
     parser.add_argument("--dp_privacy_epsilon", type=float, default=1.0)
     parser.add_argument("--dp_clip_threshold", type=float, default=0.5)
 
-    parser.add_argument("--config_path", type=str, default=None, help="Path to config file.")
-    
+    parser.add_argument(
+        "--config_path", type=str, default=None, help="Path to config file."
+    )
+
     options = parser.parse_args()
     options.agg = options.agg.lower()
     options.attack = options.attack.lower()
     if options.algorithm == "fedsgd":
         options.local_round = 1
-        
     options.dp_privacy_sensitivity = 2 * options.dp_clip_threshold / options.batch_size
     ROOT_DIR = os.path.dirname(os.path.abspath(__file__))
-    
+
+    if options.config_path:
+        options.aggregator_kws = {}
+        over_write_args_from_file(options, options.config_path)
+        options.agg = options.agg.lower()
+        options.attack = options.attack.lower()
+
+
     EXP_DIR = os.path.join(ROOT_DIR, f"outputs/{options.dataset}")
-    
+
     attack_args = {
-        'signflipping': {},
-        'noise': {"std": 0.1},
-        'labelflipping': {},
-        'permutation': {},
-        'attackclippedclustering': {},
-        'fangattack': {},
-        'distancemaximization': {},
-        'ipm': {"epsilon": 0.5},
-        'alie': {"num_clients": options.num_clients, "num_byzantine": options.num_byzantine},
+        "signflipping": {},
+        "noise": {"std": 0.1},
+        "labelflipping": {},
+        "permutation": {},
+        "attackclippedclustering": {},
+        "fangattack": {},
+        "distancemaximization": {},
+        "ipm": {"epsilon": 0.5},
+        "alie": {
+            "num_clients": options.num_clients,
+            "num_byzantine": options.num_byzantine,
+        },
     }
     agg_args = {
-        'trimmedmean': {"num_excluded": options.num_byzantine},
-        'median': {},
-        'mean': {},
-        'signguard': {},
-        'geomed': {},
-        'dnc': {"num_byzantine": options.num_byzantine},
-        'autogm': {"lamb": 2.0},
-        'clippedclustering': {"max_tau": 2.0, "signguard": True, "linkage": 'average'},
-        'clustering': {},
-        'centeredclipping': {},
-        'multikrum': {"num_excluded": options.num_byzantine, "k": 5},
+        "trimmedmean": {"num_excluded": options.num_byzantine},
+        "median": {},
+        "mean": {},
+        "signguard": {},
+        "geomed": {},
+        "dnc": {"num_byzantine": options.num_byzantine},
+        "autogm": {"lamb": 2.0},
+        "clippedclustering": {"max_tau": 2.0, "signguard": True, "linkage": "average"},
+        "clustering": {},
+        "centeredclipping": {},
+        "multikrum": {"num_excluded": options.num_byzantine, "k": 5},
     }
-    
+
     options.attack_kws = attack_args[options.attack]
     options.aggregator_kws = agg_args[options.agg]
-    
-    
 
     options.adversary_args = {
-        'fangattack': {"num_byzantine": options.num_byzantine, "agg": "median"},
-        'distancemaximization': {"num_byzantine": options.num_byzantine, "agg": "trimmedmean"},
+        "fangattack": {"num_byzantine": options.num_byzantine, "agg": "median"},
+        "distancemaximization": {
+            "num_byzantine": options.num_byzantine,
+            "agg": "trimmedmean",
+        },
     }
-    
-    if options.config_path:
-        options.aggregator_kws = {}
-        over_write_args_from_file(options, options.config_path)
-        options.agg = options.agg.lower()
-        options.attack = options.attack.lower()
-    
+
     options.log_dir = (
-            EXP_DIR
-            + f"_{options.algorithm}"
-            + f"/b{options.num_byzantine}"
-            + f"_{options.attack}" + (
-                "_" + "_".join([k + str(v) for k, v in options.attack_kws.items()]) if options.attack_kws else "")
-            + f"_{options.agg}" + (
-                "_" + "_".join([k + str(v) for k, v in options.aggregator_kws.items()]) if options.aggregator_kws else "")
-            + (f"_lr{options.lr}")
-            + (f"_serv_momentum{options.serv_momentum}")
-            + (f"_bz{options.batch_size}")
-            + (f"_localround{options.local_round}")
-            + ("_noniid" if options.non_iid else "")
-            + (f"_privacy_epsilon{options.dp_privacy_epsilon}_clip_threshold{options.dp_clip_threshold}" if options.dp else "")
-            + f"_seed{options.seed}"
+        EXP_DIR
+        + f"_{options.algorithm}"
+        + f"/b{options.num_byzantine}"
+        + f"_{options.attack}"
+        + (
+            "_" + "_".join([k + str(v) for k, v in options.attack_kws.items()])
+            if options.attack_kws
+            else ""
+        )
+        + f"_{options.agg}"
+        + (
+            "_" + "_".join([k + str(v) for k, v in options.aggregator_kws.items()])
+            if options.aggregator_kws
+            else ""
+        )
+        + (f"_lr{options.lr}")
+        + (f"_serv_momentum{options.serv_momentum}")
+        + (f"_bz{options.batch_size}")
+        + (f"_localround{options.local_round}")
+        + ("_noniid" if options.non_iid else "")
+        + (
+            f"_privacy_epsilon{options.dp_privacy_epsilon}_clip_threshold"
+            f"{options.dp_clip_threshold}"
+            if options.dp
+            else ""
+        )
+        + f"_seed{options.seed}"
     )
-    
+
     if not torch.cuda.is_available():
-        print('Unfortunaly, we currently do not have any GPU on your machine. ')
+        print("We currently do not have any GPU on your machine. ")
         options.num_gpus = 0
         options.gpu_per_actor = 0
     return options
diff --git a/scripts/main.py b/scripts/main.py
index 446206a..b80ddbc 100644
--- a/scripts/main.py
+++ b/scripts/main.py
@@ -1,22 +1,20 @@
+import math
 import os
 
 import ray
 import torch
-import os
-from blades.models.mnist import MLP
 from args import options
-from blades.simulator import Simulator
-from blades.datasets import CIFAR10
-from blades.datasets import MNIST
-import math
-from blades.models.cifar10 import CCTNet
 
+from blades.core.simulator import Simulator
+from blades.datasets import CIFAR10, CIFAR100, MNIST
+from blades.models.cifar10 import CCTNet
+from blades.models.mnist import MLP
 
 args = options
 if not ray.is_initialized():
     ray.init()
 else:
-    ray.init(address='auto')
+    ray.init(address="auto")
 
 if not os.path.exists(options.log_dir):
     os.makedirs(options.log_dir)
@@ -24,34 +22,56 @@
 # Use absolute path name, as Ray might have trouble in fetching relative path
 data_root = os.path.abspath("./data")
 
-cache_name = options.dataset + "_" + options.algorithm + ("_noniid" if not options.non_iid else "") + f"_{str(options.num_clients)}_{str(options.seed)}"
-if options.dataset == 'cifar10':
-    dataset = CIFAR10(data_root=data_root,
-                      cache_name=cache_name,
-                      train_bs=options.batch_size,
-                      num_clients=options.num_clients,
-                      iid=not options.non_iid,
-                      seed=0,
-                      )  # built-in federated cifar10 dataset
+cache_name = (
+    options.dataset
+    + "_"
+    + options.algorithm
+    + ("_noniid" if not options.non_iid else "")
+    + f"_{str(options.num_clients)}_{str(options.seed)}"
+)
+if options.dataset == "cifar10":
+    dataset = CIFAR10(
+        data_root=data_root,
+        cache_name=cache_name,
+        train_bs=options.batch_size,
+        num_clients=options.num_clients,
+        iid=not options.non_iid,
+        seed=0,
+    )  # built-in federated cifar10 dataset
     model = CCTNet()
-elif options.dataset == 'mnist':
-    dataset = MNIST(data_root=data_root,
-                    cache_name=cache_name,
-                    train_bs=options.batch_size,
-                    num_clients=options.num_clients,
-                    iid=not options.non_iid,
-                    seed=0,
-                    )  # built-in federated MNIST dataset
+elif options.dataset == "mnist":
+    dataset = MNIST(
+        data_root=data_root,
+        cache_name=cache_name,
+        train_bs=options.batch_size,
+        num_clients=options.num_clients,
+        iid=not options.non_iid,
+        seed=0,
+    )  # built-in federated MNIST dataset
     model = MLP()
+elif options.dataset == "cifar100":
+    dataset = CIFAR100(
+        data_root=data_root,
+        cache_name=cache_name,
+        train_bs=options.batch_size,
+        num_clients=options.num_clients,
+        iid=not options.non_iid,
+        seed=0,
+    )  # built-in federated cifar100 dataset
+    model = CCTNet()
+
 else:
     raise NotImplementedError
 
 if options.gpu_per_actor > 0.0:
     model = model.to("cuda")
 
+privacy_factor = (
+    args.dp_privacy_sensitivity
+    * math.sqrt(2 * math.log(1.25 / args.dp_privacy_delta))
+    / args.dp_privacy_epsilon
+)
 
-privacy_factor = args.dp_privacy_sensitivity * math.sqrt(2 * math.log(1.25 / args.dp_privacy_delta)) / args.dp_privacy_epsilon
-   
 # configuration parameters
 conf_args = {
     "dataset": dataset,
@@ -61,8 +81,9 @@
     "use_cuda": options.gpu_per_actor > 0.0,
     "attack": options.attack,  # attack strategy
     "attack_kws": options.attack_kws,
-    # "attack_kws": options.attack_args[options.attack],
-    "adversary_kws": options.adversary_args[options.attack] if options.attack in options.adversary_args else {},
+    "adversary_kws": options.adversary_args[options.attack]
+    if options.attack in options.adversary_args
+    else {},
     "num_actors": options.num_actors,  # number of training actors
     "gpu_per_actor": options.gpu_per_actor,
     "log_path": options.log_dir,
@@ -74,45 +95,53 @@
 if options.trusted_id:
     simulator.set_trusted_clients([options.trusted_id])
 
-if options.algorithm == 'fedsgd':
+if options.algorithm == "fedsgd":
     opt = torch.optim.SGD(model.parameters(), lr=0.1, momentum=options.serv_momentum)
     lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
         opt, milestones=[2000, 3000, 5000], gamma=0.5
     )
-
-    assert options.local_round == 1, f"fedsgd requires that only one SGD is taken."
+    assert options.local_round == 1, "fedsgd requires that only one SGD is taken."
 
     # runtime parameters
     run_args = {
-        "model": model,  # global model
-        "client_optimizer": 'SGD',  # server_opt, server optimizer
+        "global_model": model,  # global global_model
+        "client_optimizer": "SGD",  # server_opt, server optimizer
         "server_optimizer": opt,  # client optimizer
         "loss": "crossentropy",  # loss funcstion
         "global_rounds": options.global_round,  # number of global rounds
-        "local_steps": options.local_round,  # number of seps "client_lr": 0.1,  # learning rateteps per round
+        "local_steps": options.local_round,
         "client_lr": 1.0,
         "validate_interval": options.validate_interval,
         "server_lr_scheduler": lr_scheduler,
-        "dp_kws": {"clip_threshold": options.dp_clip_threshold, "noise_factor": privacy_factor} if options.dp 
-                    else {}
+        "dp_kws": {
+            "clip_threshold": options.dp_clip_threshold,
+            "noise_factor": privacy_factor,
+        }
+        if options.dp
+        else {},
     }
 
-
-elif options.algorithm == 'fedavg':
+elif options.algorithm == "fedavg":
     opt = torch.optim.SGD(model.parameters(), lr=0.1)
     lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
-        opt, milestones=[options.global_round / 3, options.global_round / 2, 2 * options.global_round / 3], gamma=0.5
+        opt,
+        milestones=[
+            options.global_round / 3,
+            options.global_round / 2,
+            2 * options.global_round / 3,
+        ],
+        gamma=0.5,
     )
     server_opt = torch.optim.SGD(model.parameters(), lr=1.0)
     # runtime parameters
     run_args = {
-        "model": model,  # global model
+        "global_model": model,  # global global_model
         "server_optimizer": server_opt,  # server_opt, server optimizer
         "client_optimizer": opt,  # client optimizer
         "loss": "crossentropy",  # loss funcstion
         "global_rounds": options.global_round,  # number of global rounds
-        "local_steps": options.local_round,  # number of seps "client_lr": 0.1,  # learning rateteps per round
-        # "server_lr": 1.0,
+        "local_steps": options.local_round,  # number of seps
+        "client_lr": 0.1,  # learning rateteps per round
         "validate_interval": 20,
         "client_lr_scheduler": lr_scheduler,
     }
diff --git a/setup.cfg b/setup.cfg
new file mode 100644
index 0000000..d872bdf
--- /dev/null
+++ b/setup.cfg
@@ -0,0 +1,25 @@
+[flake8]
+exclude =
+    src/blades/models/backbones,
+    docs,
+    .git,
+    __pycache__,
+    src/blades/examples,
+    build,
+    dist,
+    src/setup.py,
+    src/blades/models/backbones/cctnets/__init__.py,
+    src/blades/models/backbones/cctnets/cct.py,
+    src/blades/models/backbones/cctnets/utils/stochastic_depth.py
+
+max-line-length = 88
+max-complexity = 100
+ignore = W293,W504,W503,E203
+
+[tool.black]
+line-length = 88
+;include = '\.pyi?$'
+exclude = .git,__pycache__,src/blades/models/backbones,src/blades/examples,docs,build,dist
+
+[codespell]
+skip = ./src/3rdParty
diff --git a/src/blades/actor.py b/src/blades/actor.py
deleted file mode 100644
index 577bf98..0000000
--- a/src/blades/actor.py
+++ /dev/null
@@ -1,60 +0,0 @@
-import ray
-
-from blades.datasets.dataset import FLDataset
-
-
-@ray.remote
-class _RayActor(object):
-    """Ray Actor"""
-    
-    def __init__(self, dataset: object, *args, **kwargs):
-        """
-       Args:
-           aggregator (callable): A callable which takes a list of tensors and returns
-               an aggregated tensor.
-           log_interval (int): Control the frequency of logging training batches
-           metrics (dict): dict of metric names and their functions
-           use_cuda (bool): Use cuda or not
-           debug (bool):
-       """
-        traindls, testdls = dataset.get_dls()
-        self.dataset = FLDataset(traindls, testdls)
-    
-    def local_training(self, clients, model, local_round, lr, *args, **kwargs):
-        if "dp" in kwargs and kwargs["dp"] == True:
-            assert "clip_threshold" in kwargs and "noise_factor" in kwargs
-            dp = True
-            clip_threshold = kwargs["clip_threshold"]
-            noise_factor = kwargs["noise_factor"]
-        else:
-            dp = False
-            clip_threshold = 0
-            noise_factor = 0
-        update = []
-        
-        for i in range(len(clients)):
-            clients[i].set_para(model)
-            clients[i].set_lr(lr)
-            clients[i].on_train_round_begin()
-            data = self.dataset.get_train_data(clients[i].id(), local_round)
-            
-            clients[i].local_training(data_batches=data)
-            
-            clients[i].on_train_round_end(dp=dp, clip_threshold=clip_threshold, noise_factor=noise_factor)
-            update.append(clients[i].get_update())
-        return update
-    
-    def evaluate(self, clients, model, round_number, batch_size, metrics):
-        update = []
-        for i in range(len(clients)):
-            clients[i].set_para(model)
-            data = self.dataset.get_all_test_data(clients[i].id())
-            result = clients[i].evaluate(
-                round_number=round_number,
-                test_set=data,
-                batch_size=batch_size,
-                metrics=metrics,
-            )
-            update.append(result)
-        return update
-
diff --git a/src/blades/aggregators/__init__.py b/src/blades/aggregators/__init__.py
index 8624de0..952b124 100644
--- a/src/blades/aggregators/__init__.py
+++ b/src/blades/aggregators/__init__.py
@@ -1,18 +1,25 @@
 from .autogm import Autogm
 from .clippedclustering import Clippedclustering
 from .clustering import Clustering
+from .dnc import Dnc
+from .fltrust import Fltrust
 from .geomed import Geomed
-from .multikrum import Multikrum
 from .mean import Mean
 from .median import Median
+from .multikrum import Multikrum
+from .signguard import Signguard
 from .trimmedmean import Trimmedmean
 
-__all__ = ['Multikrum',
-           'Median',
-           'Geomed',
-           'Autogm',
-           'Mean',
-           'Clustering',
-           'Trimmedmean',
-           'Clippedclustering',
-           ]
+__all__ = [
+    "Autogm",
+    "Clippedclustering",
+    "Clustering",
+    "Dnc",
+    "Geomed",
+    "Fltrust",
+    "Mean",
+    "Median",
+    "Multikrum",
+    "Trimmedmean",
+    "Signguard",
+]
diff --git a/src/blades/aggregators/autogm.py b/src/blades/aggregators/autogm.py
index 1155967..aaa14e4 100644
--- a/src/blades/aggregators/autogm.py
+++ b/src/blades/aggregators/autogm.py
@@ -1,9 +1,10 @@
-from typing import Union, Optional, List
+from typing import List, Optional, Union
 
 import numpy as np
 import torch
 
-from blades.client import BladesClient
+from blades.core.client import BladesClient
+
 from .geomed import Geomed
 from .mean import _BaseAggregator
 
@@ -13,29 +14,45 @@ def _compute_euclidean_distance(v1, v2):
 
 
 class Autogm(_BaseAggregator):
-    r"""
-        A robust aggregator from paper `"Byzantine-Robust Aggregation in Federated Learning Empowered Industrial IoT" <https://ieeexplore.ieee.org/abstract/document/9614992>`_
-        
-        :param maxiter: Maximum number of Weiszfeld iterations. Default 100
-        :param eps: Smallest allowed value of denominator, to avoid divide by zero. Equivalently, this is a smoothing parameter. Default 1e-6.
-        :param ftol: If objective value does not improve by at least this `ftol` fraction, terminate the algorithm. Default 1e-10.
+    r"""A robust aggregator from paper `"Byzantine-Robust Aggregation in
+    Federated Learning Empowered Industrial IoT".
+
+    <https://ieeexplore.ieee.org/abstract/document/9614992>`_
+
+    Args:
+        maxiter: Maximum number of Weiszfeld iterations, default 2.0.
+        eps: Smallest allowed value of denominator, to avoid divide by zero.
+        ftol: If objective value does not improve by at least this `ftol` fraction,
+                terminate the algorithm. Default 1e-10.
     """
-    
-    def __init__(self, lamb: Optional[float] = None, maxiter: Optional[int] = 100, eps: Optional[float] = 1e-6,
-                 ftol: Optional[float] = 1e-10):
+
+    def __init__(
+        self,
+        lamb: Optional[float] = None,
+        maxiter: Optional[int] = 2.0,
+        eps: Optional[float] = 1e-6,
+        ftol: Optional[float] = 1e-10,
+    ):
         super(Autogm, self).__init__()
         self.lamb = lamb
         self.maxiter = maxiter
         self.eps = eps
         self.ftol = ftol
         self.gm_agg = Geomed(maxiter=maxiter, eps=eps, ftol=ftol)
-    
+
     def geometric_median_objective(self, median, points, alphas):
-        return sum([alpha * _compute_euclidean_distance(median, p) for alpha, p in zip(alphas, points)])
-    
-    def __call__(self, inputs: Union[List[BladesClient], List[torch.Tensor]], weights=None):
+        return sum(
+            [
+                alpha * _compute_euclidean_distance(median, p)
+                for alpha, p in zip(alphas, points)
+            ]
+        )
+
+    def __call__(
+        self, inputs: Union[List[BladesClient], List[torch.Tensor]], weights=None
+    ):
         updates = self._get_updates(inputs)
-        
+
         lamb = 1 * len(updates) if self.lamb is None else self.lamb
         alpha = np.ones(len(updates)) / len(updates)
         median = self.gm_agg(updates, alpha)
@@ -46,17 +63,17 @@ def __call__(self, inputs: Union[List[BladesClient], List[torch.Tensor]], weight
             prev_global_obj = global_obj
             for idx, local_model in enumerate(updates):
                 distance[idx] = _compute_euclidean_distance(local_model, median)
-            
+
             idxs = [x for x, _ in sorted(enumerate(distance), key=lambda x: x)]
             eta_optimal = 10000000000000000.0
             for p in range(0, len(idxs)):
-                eta = (sum([distance[i] for i in idxs[:p + 1]]) + lamb) / (p + 1)
+                eta = (sum([distance[i] for i in idxs[: p + 1]]) + lamb) / (p + 1)
                 if p < len(idxs) and eta - distance[idxs[p]] < 0:
                     break
                 else:
                     eta_optimal = eta
             alpha = np.array([max(eta_optimal - d, 0) / lamb for d in distance])
-            
+
             median = self.gm_agg(updates, alpha)
             gm_sum = self.geometric_median_objective(median, updates, alpha)
             global_obj = gm_sum + lamb * np.linalg.norm(alpha) ** 2 / 2
diff --git a/src/blades/aggregators/byzantinesgd.py b/src/blades/aggregators/byzantinesgd.py
deleted file mode 100644
index 38495cf..0000000
--- a/src/blades/aggregators/byzantinesgd.py
+++ /dev/null
@@ -1,80 +0,0 @@
-import statistics
-
-import torch
-
-from .mean import _BaseAggregator
-
-
-def _get_vectorized_parameters(optimizer):
-    c = []
-    for group in optimizer.param_groups:
-        for p in group["params"]:
-            c.append(p.data.view(-1))
-    return torch.cat(c)
-
-
-class ByzantineSGD(_BaseAggregator):
-    def __init__(self, m, th_A, th_B, th_V, optimizer):
-        # m is the number of workers
-        self.m = m
-        self.th_A = th_A
-        self.th_B = th_B
-        self.th_V = th_V
-        
-        # TODO:vectorize
-        self.optimizer = optimizer
-        
-        # TODO: vectorize the model
-        self.init_model = _get_vectorized_parameters(optimizer).clone()
-        
-        self.A = [0] * m
-        self.B = [0] * m
-        self.good = list(range(self.m))
-        self.debug_message = ""
-    
-    def vector_median(self, vs, threshold):
-        for i in range(self.m):
-            count = 0
-            for j in range(self.m):
-                distance = (vs[i] - vs[j]).norm()
-                count += distance <= threshold
-                self.debug_message += f"(i, j, d)=({i}, {j}, {distance.item():.3f})\n"
-                if count > self.m / 2:
-                    return i, vs[i]
-        raise NotImplementedError("No median found")
-    
-    def __call__(self, inputs):
-        model_diff = _get_vectorized_parameters(self.optimizer) - self.init_model
-        for i in range(self.m):
-            self.A[i] += inputs[i].dot(model_diff)
-            self.B[i] += inputs[i]
-        self.debug_message = "\n=== Aggregation begins ===\n"
-        self.debug_message += f"=> model_diff={model_diff.norm().item():.3f} \n"
-        self.debug_message += (
-                "=> A=[" + (" ".join([f"{a.item():.3f}" for a in self.A])) + "]\n"
-        )
-        
-        A_med = statistics.median(self.A)
-        self.debug_message += f"=> A_med={A_med.item():.3f}\n"
-        
-        self.debug_message += "=> Find median of B\n"
-        B_med_index, B_med = self.vector_median(self.B, self.th_B)
-        self.debug_message += f"=> Select median from worker {B_med_index}\n"
-        
-        self.debug_message += "=> Find median of current gradients\n"
-        Grad_med_index, grad_median = self.vector_median(inputs, 2 * self.th_V)
-        self.debug_message += f"=> Select median from worker {Grad_med_index}\n"
-        
-        self.debug_message += "\n=> start filtering\n"
-        candidate = []
-        for i in self.good:
-            ai2med = abs(self.A[i] - A_med).item()
-            bi2med = (self.B[i] - B_med).norm().item()
-            g2med = (inputs[i] - grad_median).norm().item()
-            self.debug_message += f"i={i} (a2m, th)=({ai2med:.3f}, {self.th_A}) (b2m, th)=({bi2med:.3f}, {self.th_B}) (g2m, th)=({g2med:.3f},{4 * self.th_V})\n"
-            if all([ai2med <= self.th_A, bi2med <= self.th_B, g2med <= 4 * self.th_V]):
-                candidate.append(i)
-        
-        self.good = candidate
-        
-        return sum(inputs[i] for i in self.good) / len(self.good)
diff --git a/src/blades/aggregators/centeredclipping.py b/src/blades/aggregators/centeredclipping.py
index 45fa28d..6194a96 100644
--- a/src/blades/aggregators/centeredclipping.py
+++ b/src/blades/aggregators/centeredclipping.py
@@ -4,47 +4,48 @@
 
 import torch
 
-from .mean import _BaseAggregator
-from .mean import _BaseAsyncAggregator
+from .mean import _BaseAggregator, _BaseAsyncAggregator
 
 debug_logger = logging.getLogger("debug")
 
 
 class Centeredclipping(_BaseAggregator):
-    r"""
-
-      A robust aggregator from paper `Learning from History for Byzantine Robust Optimization <http://proceedings.mlr.press/v139/karimireddy21a.html>`_
-      
-      It iteratively clips the updates around the center while updating the center accordingly.
-      
-      :param tau: The threshold of clipping. Default 10.0
-      :param n_iter: The number of clipping iterations. Default 5
+    r"""A robust aggregator from paper `Learning from History for Byzantine
+    Robust Optimization.
+
+    <http://proceedings.mlr.press/v139/karimireddy21a.html>`_
+
+    It iteratively clips the updates around the center while updating the center
+    accordingly.
+
+    :param tau: The threshold of clipping. Default 10.0
+    :param n_iter: The number of clipping iterations. Default 5
     """
-    
+
     def __init__(self, tau: Optional[float] = 10.0, n_iter: Optional[int] = 5):
         self.tau = tau
         self.n_iter = n_iter
         super(Centeredclipping, self).__init__()
         self.momentum = None
-    
+
     def clip(self, v):
         v_norm = torch.norm(v)
         scale = min(1, self.tau / v_norm)
         return v * scale
-    
+
     def __call__(self, clients):
         updates = list(map(lambda w: w.get_update(), clients))
         if self.momentum is None:
             self.momentum = torch.zeros_like(updates[0])
-        
+
         for _ in range(self.n_iter):
             self.momentum = (
-                    sum(self.clip(v - self.momentum) for v in updates) / len(updates)
-                    + self.momentum
+                sum(self.clip(v - self.momentum) for v in updates) / len(updates)
+                + self.momentum
             )
-        
+
         return torch.clone(self.momentum).detach()
-    
+
     def __str__(self):
         return "Clipping (tau={}, n_iter={})".format(self.tau, self.n_iter)
 
@@ -58,24 +59,24 @@ def __init__(self, node, weights, opt, model, tau, n_iter=1):
         assert len(weights.shape) == 1
         self.node = node
         self.weights = weights
-    
+
     def _vectorize_model(self, model):
         """"""
         state_dict = model.state_dict()
         return torch.cat([state_dict[k].data.view(-1) for k in state_dict])
-    
+
     def _wrap_step(self, opt: torch.optim.Optimizer, model: torch.nn.Module):
         """Wrap the step function of opt to track the change."""
-        debug_logger.info(f"Wrap the step function of opt")
-        
+        debug_logger.info("Wrap the step function of opt")
+
         if hasattr(opt, "_core_step") or hasattr(opt, "anchorclipping"):
             raise NotImplementedError
-        
+
         # Cache the old class
         opt._core_step = types.MethodType(opt.__class__.step, opt)
         opt.anchorclipping = self
-        
-        # Update the anchor vector y everytime the `step` is called.
+
+        # Update the anchor vector y every time the `step` is called.
         def anchor_clipping_step(self, closure=None):
             state_dict = model.state_dict()
             flattened = torch.cat([state_dict[k].data.view(-1) for k in state_dict])
@@ -86,9 +87,9 @@ def anchor_clipping_step(self, closure=None):
             )
             # debug_logger.info((after_flattened - flattened)[:5])
             self.anchorclipping._anchor_buffer.add_(after_flattened - flattened)
-        
+
         opt.step = types.MethodType(anchor_clipping_step, opt)
-    
+
     def __call__(self, inputs):
         assert len(inputs) == 1 + len(self.node.edges)
         clipped = self._anchor_buffer + self.clip(inputs[0] - self._anchor_buffer)
@@ -98,40 +99,40 @@ def __call__(self, inputs):
             clipped = self._anchor_buffer + self.clip(inp - self._anchor_buffer)
             s += self.weights[theothernode.index] * clipped
         return s
-    
+
     def __str__(self):
         return "_AnchorClipping(tau={}, n_iter={})".format(self.tau, self.n_iter)
 
 
 class _AsyncCenteredClipping(_BaseAsyncAggregator):
-    """
-    Comparing to Clipping, _AsyncCenteredClipping does not average the clipped gradient but use
-    fraction $1 / n$ where `n` is the number of total gradients.
+    """Comparing to Clipping, _AsyncCenteredClipping does not average the
+    clipped gradient but use fraction $1 / n$ where `n` is the number of total
+    gradients."""
 
-    """
-    
     def __init__(self, tau, n_iter=1):
         self.tau = tau
         self.n_iter = n_iter
         super(_AsyncCenteredClipping, self).__init__()
-        
+
         self.momentum = 0
-    
+
     def clip(self, v):
         v_norm = torch.norm(v)
         scale = min(1, self.tau / v_norm)
         return v * scale
-    
+
     def __call__(self, inputs):
         n = len(inputs)
         filtered = list(filter(lambda x: x is not None, inputs))
-        
+
         for _ in range(self.n_iter):
             self.momentum = (
-                    sum(self.clip(v - self.momentum) for v in filtered) / n + self.momentum
+                sum(self.clip(v - self.momentum) for v in filtered) / n + self.momentum
             )
-        
+
         return torch.clone(self.momentum).detach()
-    
+
     def __str__(self):
-        return "_AsyncCenteredClipping (tau={}, n_iter={})".format(self.tau, self.n_iter)
+        return "_AsyncCenteredClipping (tau={}, n_iter={})".format(
+            self.tau, self.n_iter
+        )
diff --git a/src/blades/aggregators/clippedclustering.py b/src/blades/aggregators/clippedclustering.py
index e187596..2472d17 100644
--- a/src/blades/aggregators/clippedclustering.py
+++ b/src/blades/aggregators/clippedclustering.py
@@ -1,69 +1,84 @@
-from typing import Union, List
+from typing import List, Union
 
 import numpy as np
 import torch
 from numpy import inf
 from sklearn.cluster import AgglomerativeClustering, KMeans
 
-from blades.client import BladesClient
+from blades.core.client import BladesClient
 from blades.utils import torch_utils
-from .mean import _BaseAggregator, Mean
+
+from .mean import Mean, _BaseAggregator
 from .median import Median
 
 
 class Clippedclustering(_BaseAggregator):
-    r"""
-         A robust aggregator from paper `"An Experimental Study of Byzantine-Robust sAggregation Schemes in Federated Learning" <https://www.techrxiv.org/articles/preprint/An_Experimental_Study_of_Byzantine-Robust_Aggregation_Schemes_in_Federated_Learning/19560325>`_
+    r"""A robust aggregator from paper `"An Experimental Study of Byzantine-
+    Robust sAggregation Schemes in Federated Learning" <https://www.techrxiv.or
+    g/articles/preprint/An_Experimental_Study_of_Byzantine.
+
+         -Robust_Aggregation_Schemes_in_Federated_Learning/19560325>`_ # noqa: E501
 
          it separates the client population into two groups based on the cosine similarities
 
     Args:
-        tau (float): threshold of clipping norm. If it is not given, updates are clipped according the median of historical norm.
+        tau (float): threshold of clipping norm.
+                    If it is not given, updates are clipped according the median of historical norm.
     """
-    
-    def __init__(self, agg='mean', signguard=False, max_tau=1e5, linkage='average') -> None:
+
+    def __init__(
+        self, agg="mean", signguard=False, max_tau=1e5, linkage="average"
+    ) -> None:
         super(Clippedclustering, self).__init__()
-        
-        assert linkage in ['average', 'single']
+
+        assert linkage in ["average", "single"]
         self.tau = max_tau
         self.signguard = signguard
         self.linkage = linkage
         self.l2norm_his = []
-        if agg == 'mean':
+        if agg == "mean":
             self.agg = Mean()
-        elif agg == 'median':
+        elif agg == "median":
             self.agg = Median()
         else:
             raise NotImplementedError(f"{agg} is not supported yet.")
-    
-    def __call__(self, inputs: Union[List[BladesClient], List[torch.Tensor], torch.Tensor]):
+
+    def __call__(
+        self, inputs: Union[List[BladesClient], List[torch.Tensor], torch.Tensor]
+    ):
         updates = self._get_updates(inputs)
         l2norms = [torch.norm(update).item() for update in updates]
         self.l2norm_his.extend(l2norms)
         threshold = np.median(self.l2norm_his)
         threshold = min(threshold, self.tau)
-        
+
         # print(threshold, l2norms)
         for idx, l2 in enumerate(l2norms):
             if l2 > threshold:
                 updates[idx] = torch_utils.clip_tensor_norm_(updates[idx], threshold)
-        
+
         num = len(updates)
         dis_max = np.zeros((num, num))
         for i in range(num):
-            for j in range(i+1, num):
-                dis_max[i, j] = 1 - torch.nn.functional.cosine_similarity(updates[i, :], updates[j, :], dim=0)
-                dis_max[j, i] = dis_max[i, j] 
+            for j in range(i + 1, num):
+                dis_max[i, j] = 1 - torch.nn.functional.cosine_similarity(
+                    updates[i, :], updates[j, :], dim=0
+                )
+                dis_max[j, i] = dis_max[i, j]
         dis_max[dis_max == -inf] = 0
         dis_max[dis_max == inf] = 2
         dis_max[np.isnan(dis_max)] = 2
-        clustering = AgglomerativeClustering(affinity='precomputed', linkage=self.linkage, n_clusters=2)
+        clustering = AgglomerativeClustering(
+            affinity="precomputed", linkage=self.linkage, n_clusters=2
+        )
         clustering.fit(dis_max)
 
         flag = 1 if np.sum(clustering.labels_) > num // 2 else 0
-        S1_idxs = list([idx for idx, label in enumerate(clustering.labels_) if label == flag])
+        S1_idxs = list(
+            [idx for idx, label in enumerate(clustering.labels_) if label == flag]
+        )
         selected_idxs = S1_idxs
-        
+
         if self.signguard:
             features = []
             num_para = len(updates[0])
@@ -71,19 +86,21 @@ def __call__(self, inputs: Union[List[BladesClient], List[torch.Tensor], torch.T
                 feature0 = (update > 0).sum().item() / num_para
                 feature1 = (update < 0).sum().item() / num_para
                 feature2 = (update == 0).sum().item() / num_para
-        
+
                 features.append([feature0, feature1, feature2])
-    
+
             kmeans = KMeans(n_clusters=2, random_state=0).fit(features)
-    
+
             flag = 1 if np.sum(kmeans.labels_) > num // 2 else 0
-            S2_idxs = list([idx for idx, label in enumerate(kmeans.labels_) if label == flag])
+            S2_idxs = list(
+                [idx for idx, label in enumerate(kmeans.labels_) if label == flag]
+            )
 
             selected_idxs = list(set(S1_idxs) & set(S2_idxs))
-        
+
         benign_updates = []
         for idx in selected_idxs:
             benign_updates.append(updates[idx])
-                
+
         values = self.agg(torch.vstack(benign_updates))
         return values
diff --git a/src/blades/aggregators/clustering.py b/src/blades/aggregators/clustering.py
index 73d1e95..f0be36b 100644
--- a/src/blades/aggregators/clustering.py
+++ b/src/blades/aggregators/clustering.py
@@ -1,40 +1,54 @@
-from typing import Union, List
+from typing import List, Union
 
 import numpy as np
 import torch
 from numpy import inf
 from sklearn.cluster import AgglomerativeClustering
 
-from blades.client import BladesClient
+from blades.core.client import BladesClient
+
 from .mean import _BaseAggregator
 
 
 class Clustering(_BaseAggregator):
-    r"""
-     A robust aggregator from paper `"On the byzantine robustness of clustered federated learning" <https://ieeexplore.ieee.org/abstract/document/9054676>`_
-     
-     it separates the client population into two groups based on the cosine similarities
+    r"""A robust aggregator from paper `"On the byzantine robustness of
+    clustered federated learning".
+
+    <https://ieeexplore.ieee.org/abstract/document/9054676>`_
+
+    it separates the client population into two groups based on the cosine similarities
     """
-    
+
     def __init__(self):
         super(Clustering, self).__init__()
-    
-    def __call__(self, inputs: Union[List[BladesClient], List[torch.Tensor], torch.Tensor]):
+
+    def __call__(
+        self, inputs: Union[List[BladesClient], List[torch.Tensor], torch.Tensor]
+    ):
         updates = self._get_updates(inputs)
         num = len(updates)
         dis_max = np.zeros((num, num))
         for i in range(num):
-            for j in range(i+1, num):
-                dis_max[i, j] = 1 - torch.nn.functional.cosine_similarity(updates[i, :], updates[j, :], dim=0)
+            for j in range(i + 1, num):
+                dis_max[i, j] = 1 - torch.nn.functional.cosine_similarity(
+                    updates[i, :], updates[j, :], dim=0
+                )
                 dis_max[j, i] = dis_max[i, j]
         dis_max[dis_max == -inf] = -1
         dis_max[dis_max == inf] = 1
         dis_max[np.isnan(dis_max)] = -1
         # with open('../notebooks/updates_fedsgd_ipm.npy', 'wb') as f:
         #     np.save(f, dis_max)
-        clustering = AgglomerativeClustering(affinity='precomputed', linkage='complete', n_clusters=2)
+        clustering = AgglomerativeClustering(
+            affinity="precomputed", linkage="complete", n_clusters=2
+        )
         clustering.fit(dis_max)
         flag = 1 if np.sum(clustering.labels_) > num // 2 else 0
-        values = torch.vstack(list(model for model, label in zip(updates, clustering.labels_) if label == flag)).mean(
-            dim=0)
+        values = torch.vstack(
+            list(
+                model
+                for model, label in zip(updates, clustering.labels_)
+                if label == flag
+            )
+        ).mean(dim=0)
         return values
diff --git a/src/blades/aggregators/dnc.py b/src/blades/aggregators/dnc.py
index dda8021..176a48e 100644
--- a/src/blades/aggregators/dnc.py
+++ b/src/blades/aggregators/dnc.py
@@ -1,42 +1,52 @@
-from typing import Union, List
+from typing import List, Union
 
 import numpy as np
 import torch
 
-from blades.client import BladesClient
+from blades.core.client import BladesClient
+
 from .mean import _BaseAggregator
 
 
 class Dnc(_BaseAggregator):
-    r"""
-         A robust aggregator from paper `Manipulating the Byzantine: Optimizing Model Poisoning Attacks and Defenses for Federated Learning <https://par.nsf.gov/servlets/purl/10286354>`_
+    r"""A robust aggregator from paper `Manipulating the Byzantine: Optimizing
+    Model Poisoning Attacks and Defenses for Federated Learning.
 
+    <https://par.nsf.gov/servlets/purl/10286354>`_
     """
-    
-    def __init__(self, num_byzantine, *, sub_dim=10000, num_iters=1, filter_frac=1.0) -> None:
+
+    def __init__(
+        self, num_byzantine, *, sub_dim=10000, num_iters=1, filter_frac=1.0
+    ) -> None:
         super(Dnc, self).__init__()
-        
+
         self.num_byzantine = num_byzantine
         self.sub_dim = sub_dim
         self.num_iters = num_iters
         self.fliter_frac = filter_frac
-    
-    def __call__(self, inputs: Union[List[BladesClient], List[torch.Tensor], torch.Tensor]):
+
+    def __call__(
+        self, inputs: Union[List[BladesClient], List[torch.Tensor], torch.Tensor]
+    ):
         updates = self._get_updates(inputs)
         d = len(updates[0])
-                        
+
         benign_ids = []
         for i in range(self.num_iters):
-            indices = torch.randperm(d)[:self.sub_dim]
+            indices = torch.randperm(d)[: self.sub_dim]
             sub_updates = updates[:, indices]
             mu = sub_updates.mean(dim=0)
             centered_update = sub_updates - mu
             v = torch.linalg.svd(centered_update, full_matrices=False)[2][0, :]
-            s = np.array([(torch.dot(update - mu, v) ** 2).item() for update in sub_updates])
+            s = np.array(
+                [(torch.dot(update - mu, v) ** 2).item() for update in sub_updates]
+            )
 
-            good = s.argsort()[:len(updates) - int(self.fliter_frac * self.num_byzantine)]
+            good = s.argsort()[
+                : len(updates) - int(self.fliter_frac * self.num_byzantine)
+            ]
             benign_ids.extend(good)
-            
+
         benign_ids = list(set(benign_ids))
         benign_updates = updates[benign_ids, :].mean(dim=0)
         return benign_updates
diff --git a/src/blades/aggregators/fltrust.py b/src/blades/aggregators/fltrust.py
index d81e34c..7b6bce4 100644
--- a/src/blades/aggregators/fltrust.py
+++ b/src/blades/aggregators/fltrust.py
@@ -2,37 +2,43 @@
 
 import torch
 
-from blades.client import BladesClient
+from blades.core.client import BladesClient
 
 
 class Fltrust(object):
     r"""
-         ``Fltrust`` it a trusted-based aggregator from paper `FLTrust: Byzantine-robust Federated Learning via Trust Bootstrapping <https://arxiv.org/abs/2012.13995>`_
+    ``Fltrust`` it a trusted-based aggregator from paper
+    `FLTrust: Byzantine-robust Federated Learning via Trust Bootstrapping
+     <https://arxiv.org/abs/2012.13995>`_
+    """
 
-         
-        """
-    
     def __call__(self, clients: List[BladesClient]):
         trusted_clients = [client for client in clients if client.is_trusted()]
         assert len(trusted_clients) == 1
         trusted_client = trusted_clients[0]
-        
+
         untrusted_clients = [client for client in clients if not client.is_trusted()]
         trusted_update = trusted_client.get_update()
         trusted_norm = torch.norm(trusted_update).item()
         untrusted_updates = list(map(lambda w: w.get_update(), untrusted_clients))
         cosine_similarity = torch.nn.CosineSimilarity(dim=0, eps=1e-6)
-        ts = torch.Tensor(list(map(lambda update:
-                                   torch.nn.functional.relu(
-                                       cosine_similarity(trusted_update, update)
-                                   ),
-                                   untrusted_updates,
-                                   ))
-                          )
-        pseudo_gradients = torch.vstack(list(map(lambda update: update * trusted_norm / torch.norm(update).item(),
-                                                 untrusted_updates,
-                                                 )
-                                             )
-                                        )
+        ts = torch.Tensor(
+            list(
+                map(
+                    lambda update: torch.nn.functional.relu(
+                        cosine_similarity(trusted_update, update)
+                    ),
+                    untrusted_updates,
+                )
+            )
+        )
+        pseudo_gradients = torch.vstack(
+            list(
+                map(
+                    lambda update: update * trusted_norm / torch.norm(update).item(),
+                    untrusted_updates,
+                )
+            )
+        )
         true_update = (pseudo_gradients.T @ ts) / ts.sum()
         return true_update
diff --git a/src/blades/aggregators/geomed.py b/src/blades/aggregators/geomed.py
index 11cddb2..6650de0 100644
--- a/src/blades/aggregators/geomed.py
+++ b/src/blades/aggregators/geomed.py
@@ -1,9 +1,10 @@
-from typing import Union, Optional, List
+from typing import List, Optional, Union
 
 import numpy as np
 import torch
 
-from blades.client import BladesClient
+from blades.core.client import BladesClient
+
 from .mean import _BaseAggregator
 
 
@@ -15,16 +16,16 @@ def smoothed_weiszfeld(weights, alphas, z, eps=1e-6, T=5):
     m = len(weights)
     if len(alphas) != m:
         raise ValueError
-    
+
     if eps < 0:
         raise ValueError
-    
+
     for t in range(T):
         betas = []
         for k in range(m):
             distance = _compute_euclidean_distance(z, weights[k])
             betas.append(max(eps, alphas[k] / max(distance, eps)))
-        
+
         z = 0
         for w, beta in zip(weights, betas):
             z += w * beta
@@ -33,35 +34,51 @@ def smoothed_weiszfeld(weights, alphas, z, eps=1e-6, T=5):
 
 
 class Geomed(_BaseAggregator):
-    r"""
+    r"""A robust aggregator from paper `"Distributed Statistical Machine
+    Learning in Adversarial Settings: Byzantine Gradient Descent".
+
+    <https://arxiv.org/abs/1705.05491>`_
+
+    ``GeoMed`` aims to find a vector that minimizes the sum of its Euclidean distances
+    to all the update vectors:
 
-    A robust aggregator from paper `"Distributed Statistical Machine Learning in Adversarial Settings: Byzantine Gradient Descent" <https://arxiv.org/abs/1705.05491>`_
-    
-    ``GeoMed`` aims to find a vector that minimizes the sum of its Euclidean distances to all the update vectors:
-    
     .. math::
-        GeoMed := \arg\min_{\boldsymbol{z}}  \sum_{k \in [K]} \lVert \boldsymbol{z} -  {\Delta}_i \rVert.
-    
-    
-    There is no closed-form solution to the ``GeoMed`` problem. It is approximately solved using
-    Weiszfeld's algorithm in this implementation to.
-    
+        GeoMed := \arg\min_{\boldsymbol{z}}  \sum_{k \in [K]} \lVert \boldsymbol{z} -
+         {\Delta}_i \rVert.
+
+
+    There is no closed-form solution to the ``GeoMed`` problem. It is approximately
+    solved using Weiszfeld's algorithm in this implementation to.
+
     :param maxiter: Maximum number of Weiszfeld iterations. Default 100
     :param eps: Smallest allowed value of denominator, to avoid divide by zero.
-    	Equivalently, this is a smoothing parameter. Default 1e-6.
-    :param ftol: If objective value does not improve by at least this `ftol` fraction, terminate the algorithm. Default 1e-10.
+    Equivalently, this is a smoothing parameter. Default 1e-6.
+    :param ftol: If objective value does not improve by at least this `ftol` fraction,
+            terminate the algorithm, default 1e-10.
     """
-    
-    def __init__(self, maxiter: Optional[int] = 100, eps: Optional[float] = 1e-6, ftol: Optional[float] = 1e-10):
+
+    def __init__(
+        self,
+        maxiter: Optional[int] = 100,
+        eps: Optional[float] = 1e-6,
+        ftol: Optional[float] = 1e-10,
+    ):
         self.maxiter = maxiter
         self.eps = eps
         self.ftol = ftol
         super(Geomed, self).__init__()
-    
+
     def _geometric_median_objective(self, median, points, alphas):
-        return sum([alpha * _compute_euclidean_distance(median, p) for alpha, p in zip(alphas, points)])
-    
-    def __call__(self, inputs: Union[List[BladesClient], List[torch.Tensor]], weights=None):
+        return sum(
+            [
+                alpha * _compute_euclidean_distance(median, p)
+                for alpha, p in zip(alphas, points)
+            ]
+        )
+
+    def __call__(
+        self, inputs: Union[List[BladesClient], List[torch.Tensor]], weights=None
+    ):
         updates = self._get_updates(inputs)
         if weights is None:
             weights = np.ones(len(updates)) / len(updates)
@@ -69,16 +86,26 @@ def __call__(self, inputs: Union[List[BladesClient], List[torch.Tensor]], weight
         num_oracle_calls = 1
         obj_val = self._geometric_median_objective(median, updates, weights)
         for i in range(self.maxiter):
-            prev_median, prev_obj_val = median, obj_val
+            _, prev_obj_val = median, obj_val
             weights = np.asarray(
-                [max(self.eps, alpha / max(self.eps, _compute_euclidean_distance(median, p).item())) for alpha, p in
-                 zip(weights, updates)],
-                dtype=weights.dtype)
+                [
+                    max(
+                        self.eps,
+                        alpha
+                        / max(self.eps, _compute_euclidean_distance(median, p).item()),
+                    )
+                    for alpha, p in zip(weights, updates)
+                ],
+                dtype=weights.dtype,
+            )
             weights = weights / weights.sum()
-            median = torch.sum(torch.vstack([w * beta for w, beta in zip(updates, weights)]), dim=0)
+            median = torch.sum(
+                torch.vstack([w * beta for w, beta in zip(updates, weights)]),
+                dim=0,
+            )
             num_oracle_calls += 1
             obj_val = self._geometric_median_objective(median, updates, weights)
             if abs(prev_obj_val - obj_val) < self.ftol * obj_val:
                 break
-        
+
         return median
diff --git a/src/blades/aggregators/mean.py b/src/blades/aggregators/mean.py
index 09bb27f..09d599f 100644
--- a/src/blades/aggregators/mean.py
+++ b/src/blades/aggregators/mean.py
@@ -1,9 +1,9 @@
 import logging
-from typing import Union, List
+from typing import List, Union
 
 import torch
 
-from blades.client import BladesClient
+from blades.core.client import BladesClient
 
 
 class _BaseAggregator(object):
@@ -12,21 +12,25 @@ class _BaseAggregator(object):
     Args:
         dist_communicator (object): A link object which can broadcast / gather, etc.
     """
-    
+
     def __init__(self, *args, **kwargs):
         pass
         # log("Init aggregators: " + self.__str__())
         # log_dict({"Aggregator": self.__str__(), "Type": "Setup"})
-    
-    def _get_updates(self, inputs: Union[List[BladesClient], List[torch.Tensor], torch.Tensor]):
+
+    def _get_updates(
+        self, inputs: Union[List[BladesClient], List[torch.Tensor], torch.Tensor]
+    ):
         if all(isinstance(element, BladesClient) for element in inputs):
             updates = torch.stack(list(map(lambda w: w.get_update(), inputs)))
-        elif isinstance(inputs, List) and all(isinstance(element, torch.Tensor) for element in inputs):
+        elif isinstance(inputs, List) and all(
+            isinstance(element, torch.Tensor) for element in inputs
+        ):
             updates = torch.stack(inputs, dim=0)
         else:
             updates = inputs
         return updates
-    
+
     def __call__(self, inputs):
         """Aggregate the inputs and update in-place.
 
@@ -40,13 +44,13 @@ def __call__(self, inputs):
 
 
 class _BaseAsyncAggregator(object):
-    """AsyncAggregator base object"""
-    
+    """AsyncAggregator base object."""
+
     def __init__(self):
         pass
         # log("Init aggregators: " + self.__str__())
         # log_dict({"Aggregator": self.__str__(), "Type": "Setup"})
-    
+
     def __call__(self, inputs):
         """Aggregate the inputs and update in-place.
 
@@ -60,18 +64,18 @@ def __call__(self, inputs):
 
 
 class Mean(_BaseAggregator):
-    r"""
-    Computes the ``sample mean`` over the updates from all give clients.
-    """
-    
+    r"""Computes the ``sample mean`` over the updates from all give clients."""
+
     def __int__(self):
         super(Mean, self).__init__()
-    
-    def __call__(self, inputs: Union[List[BladesClient], List[torch.Tensor], torch.Tensor]):
+
+    def __call__(
+        self, inputs: Union[List[BladesClient], List[torch.Tensor], torch.Tensor]
+    ):
         updates = self._get_updates(inputs)
         values = updates.mean(dim=0)
         return values
-    
+
     def __str__(self):
         return "Mean"
 
@@ -81,16 +85,17 @@ def __call__(self, inputs):
         filtered = list(filter(lambda x: x is not None, inputs))
         values = torch.stack(filtered, dim=0).sum(dim=0) / len(inputs)
         return values
-    
+
     def __str__(self):
         return "_AsyncMean"
 
 
 class _DecentralizedAggregator(_BaseAggregator):
+    """This aggregators is applied to all nodes.
+
+    It has access to the node information and a row of mixing matrix.
     """
-    This aggregators is applied to all nodes. It has access to the node information and a row of mixing matrix.
-    """
-    
+
     def __init__(self, node, weights):
         super().__init__()
         assert len(weights.shape) == 1
@@ -99,11 +104,12 @@ def __init__(self, node, weights):
         logging.getLogger("debug").info(
             f"Aggregator: node={node.index} weights={weights}"
         )
-    
+
     def __call__(self, inputs):
-        """
-        The `inputs` is a list of tensors. The first element is the weight of itself, the second to the last elements are the
-        gradient of its neighbors.
+        """The `inputs` is a list of tensors.
+
+        The first element is the weight of itself, the second to the last
+        elements are the gradient of its neighbors.
         """
         assert len(inputs) == 1 + len(self.node.edges)
         s = self.weights[self.node.index] * inputs[0]
@@ -111,6 +117,6 @@ def __call__(self, inputs):
             theothernode = e.theother(self.node)
             s += self.weights[theothernode.index] * inp
         return s
-    
+
     def __str__(self):
         return "_DecentralizedAggregator"
diff --git a/src/blades/aggregators/median.py b/src/blades/aggregators/median.py
index 6b7181c..e9dc5be 100644
--- a/src/blades/aggregators/median.py
+++ b/src/blades/aggregators/median.py
@@ -1,24 +1,32 @@
-from typing import Union, List
+from typing import List, Union
 
 import torch
 
-from blades.client import BladesClient
+from blades.core.client import BladesClient
+
 from .mean import _BaseAggregator
 
 
 class Median(_BaseAggregator):
-    r"""
-    
-      A robust aggregator from paper `Byzantine-robust distributed learning: Towards optimal statistical rates <https://proceedings.mlr.press/v80/yin18a>`_
-      
-      It computes the coordinate-wise median of the given set of clients
-      
+    r"""A robust aggregator from paper `Byzantine-robust distributed learning:
+    Towards optimal statistical rates.
+
+    <https://proceedings.mlr.press/v80/yin18a>`_
+
+    It computes the coordinate-wise median of the given set of clients
     """
-    
+
     def __int__(self):
         super(Median, self).__init__()
-    
-    def __call__(self, inputs: Union[List[BladesClient], List[torch.Tensor]]):
+
+    def __call__(
+        self,
+        inputs: Union[
+            List[BladesClient],
+            List[torch.Tensor],
+            torch.Tensor,
+        ],
+    ):
         updates = self._get_updates(inputs)
         values_upper, _ = updates.median(dim=0)
         values_lower, _ = (-updates).median(dim=0)
diff --git a/src/blades/aggregators/multikrum.py b/src/blades/aggregators/multikrum.py
index a195153..2f92978 100644
--- a/src/blades/aggregators/multikrum.py
+++ b/src/blades/aggregators/multikrum.py
@@ -1,16 +1,18 @@
-from typing import Union, List
+from typing import List, Union
 
 import torch
 
-from blades.client import BladesClient
+from blades.core.client import BladesClient
+
 from .mean import _BaseAggregator
 
 
 def _compute_scores(distances, i, n, f):
     """Compute scores for node i.
 
-    Arguments:
-        distances {dict} -- A dict of dict of distance. distances[i][j] = dist. i, j starts with 0.
+    Args:
+        distances {dict} -- A dict of dict of distance. distances[i][j] = dist.
+        i, j starts with 0.
         i {int} -- index of worker, starting from 0.
         n {int} -- total number of workers
         f {int} -- Total number of Byzantine workers.
@@ -26,10 +28,11 @@ def _compute_scores(distances, i, n, f):
 
 
 def _multi_krum(distances, n, f, m):
-    """Multi_Krum algorithm
+    """Multi_Krum algorithm.
 
     Arguments:
-        distances {dict} -- A dict of dict of distance. distances[i][j] = dist. i, j starts with 0.
+        distances {dict} -- A dict of dict of distance. distances[i][j] = dist.
+         i, j starts with 0.
         n {int} -- Total number of workers.
         f {int} -- Total number of Byzantine workers.
         m {int} -- Number of workers for aggregation.
@@ -41,26 +44,25 @@ def _multi_krum(distances, n, f, m):
         raise ValueError(
             "Number of workers should be positive integer. Got {}.".format(f)
         )
-    
+
     if m < 1 or m > n:
         raise ValueError(
             "Number of workers for aggregation should be >=1 and <= {}. Got {}.".format(
                 m, n
             )
         )
-    
+
     if 2 * f + 2 > n:
         raise ValueError("Too many Byzantine workers: 2 * {} + 2 >= {}.".format(f, n))
-    
+
     for i in range(n - 1):
         for j in range(i + 1, n):
             if distances[i][j] < 0:
                 raise ValueError(
-                    "The distance between node {} and {} should be non-negative: Got {}.".format(
-                        i, j, distances[i][j]
-                    )
+                    "The distance between node {} and {} should be non-negative: "
+                    "Got {}.".format(i, j, distances[i][j])
                 )
-    
+
     scores = [(i, _compute_scores(distances, i, n, f)) for i in range(n)]
     sorted_scores = sorted(scores, key=lambda x: x[1])
     return list(map(lambda x: x[0], sorted_scores))[:m]
@@ -81,7 +83,7 @@ def _pairwise_euclidean_distances(vectors):
     """
     n = len(vectors)
     vectors = [v.flatten() for v in vectors]
-    
+
     distances = {}
     for i in range(n - 1):
         distances[i] = {}
@@ -91,37 +93,40 @@ def _pairwise_euclidean_distances(vectors):
 
 
 class Multikrum(_BaseAggregator):
-    r"""
-    
-      A robust aggregator from paper `"Machine Learning with Adversaries: Byzantine Tolerant Gradient Descent" <https://proceedings.neurips.cc/paper/2017/hash/f4b9ec30ad9f68f89b29639786cb62ef-Abstract.html>`_
-      
-      Given a collection of vectors, ``Krum`` strives to find one of the vector that is closest
-      to another :math:`K-M-2` ones with respect to squared Euclidean distance, which can be expressed by:
-    
-        .. math::
-           Krum := \{{\Delta}_i | i = \arg\min_{i \in [K]} \sum_{i \rightarrow j}  \lVert {\Delta}_i -
-                    {\Delta}_j \rVert^2 \}
-    
-      where :math:`i \rightarrow j` is the indices of the :math:`K-M-2` nearest neighbours of :math:`{\Delta}_i`
-      measured by squared ``Euclidean distance``,  :math:`K` is the number of input in total, and :math:`M`
-      is the number of Byzantine input.
-      
-      Args:
-            lr (float): target learning rate.
-            
+    r"""A robust aggregator from paper `"Machine Learning with Adversaries:
+    Byzantine Tolerant Gradient Descent".
+
+    <https://proceedings.neurips.cc/paper/2017/hash.
+
+    /f4b9ec30ad9f68f89b29639786cb62ef-Abstract.html>`_
+
+    Given a collection of vectors, ``Krum`` strives to find one of the vector that is
+    closest to another :math:`K-M-2` ones with respect to squared Euclidean distance,
+    which can be expressed by:
+
+      .. math::
+         Krum := \{{\Delta}_i | i = \arg\min_{i \in [K]} \sum_{i \rightarrow j}  \lVert
+         {\Delta}_i - {\Delta}_j \rVert^2 \}
+
+    where :math:`i \rightarrow j` is the indices of the :math:`K-M-2` nearest neighbours
+    of :math:`{\Delta}_i` measured by squared ``Euclidean distance``,  :math:`K` is the
+    number of input in total, and :math:`M` is the number of Byzantine input.
+
+    Args:
+          lr (float): target learning rate.
     """
-    
+
     def __init__(self, num_excluded=5, k=1):
         self.f = num_excluded
         self.m = k
         super(Multikrum, self).__init__()
-    
+
     def __call__(self, inputs: Union[List[BladesClient], List[torch.Tensor]]):
         updates = self._get_updates(inputs)
         distances = _pairwise_euclidean_distances(updates)
         top_m_indices = _multi_krum(distances, len(updates), self.f, self.m)
         values = torch.stack([updates[i] for i in top_m_indices], dim=0).mean(dim=0)
         return values
-    
+
     def __str__(self):
         return "Krum (m={})".format(self.m)
diff --git a/src/blades/aggregators/signguard.py b/src/blades/aggregators/signguard.py
index 1ee3262..eebf9fb 100644
--- a/src/blades/aggregators/signguard.py
+++ b/src/blades/aggregators/signguard.py
@@ -1,41 +1,40 @@
-import os
-import sys
-from typing import Union, List
+from typing import List, Union
 
 import numpy as np
 import torch
 from sklearn.cluster import KMeans
 
-from blades.client import BladesClient
-from .mean import _BaseAggregator, Mean
-from .median import Median
-
-file_dir = os.path.dirname(__file__)
-sys.path.append(file_dir)
-
+from blades.core.client import BladesClient
 from blades.utils import torch_utils
 
+from .mean import Mean, _BaseAggregator
+from .median import Median
+
 
 class Signguard(_BaseAggregator):
-    r"""
-         A robust aggregator from paper `Xu et al. SignGuard: Byzantine-robust Federated Learning through Collaborative Malicious Gradient Filtering <https://arxiv.org/abs/2109.05872>`_
+    r"""A robust aggregator from paper `Xu et al. SignGuard: Byzantine-robust
+    Federated Learning through Collaborative Malicious Gradient Filtering.
+
+    <https://arxiv.org/abs/2109.05872>`_
     """
-    
-    def __init__(self, agg='mean', max_tau=1e5, linkage='average') -> None:
+
+    def __init__(self, agg="mean", max_tau=1e5, linkage="average") -> None:
         super(Signguard, self).__init__()
-        
-        assert linkage in ['average', 'single']
+
+        assert linkage in ["average", "single"]
         self.tau = max_tau
         self.linkage = linkage
         self.l2norm_his = []
-        if agg == 'mean':
+        if agg == "mean":
             self.agg = Mean()
-        elif agg == 'median':
+        elif agg == "median":
             self.agg = Median()
         else:
             raise NotImplementedError(f"{agg} is not supported yet.")
-    
-    def __call__(self, inputs: Union[List[BladesClient], List[torch.Tensor], torch.Tensor]):
+
+    def __call__(
+        self, inputs: Union[List[BladesClient], List[torch.Tensor], torch.Tensor]
+    ):
         updates = self._get_updates(inputs)
         num = len(updates)
         l2norms = [torch.norm(update).item() for update in updates]
@@ -48,21 +47,23 @@ def __call__(self, inputs: Union[List[BladesClient], List[torch.Tensor], torch.T
             if l2norm >= L * M and l2norm <= R * M:
                 # S1.append(update)
                 S1_idxs.append(idx)
-                
+
         features = []
         num_para = len(updates[0])
         for update in updates:
             feature0 = (update > 0).sum().item() / num_para
             feature1 = (update < 0).sum().item() / num_para
             feature2 = (update == 0).sum().item() / num_para
-            
+
             features.append([feature0, feature1, feature2])
 
         kmeans = KMeans(n_clusters=2, random_state=0).fit(features)
         print(kmeans)
 
         flag = 1 if np.sum(kmeans.labels_) > num // 2 else 0
-        S2_idxs = list([idx for idx, label in enumerate(kmeans.labels_) if label==flag])
+        S2_idxs = list(
+            [idx for idx, label in enumerate(kmeans.labels_) if label == flag]
+        )
 
         inter = list(set(S1_idxs) & set(S2_idxs))
 
@@ -71,6 +72,6 @@ def __call__(self, inputs: Union[List[BladesClient], List[torch.Tensor], torch.T
             if l2norms[idx] > M:
                 updates[idx] = torch_utils.clip_tensor_norm_(updates[idx], M)
             benign_updates.append(updates[idx])
-        
+
         values = self.agg(torch.vstack(benign_updates))
         return values
diff --git a/src/blades/aggregators/trimmedmean.py b/src/blades/aggregators/trimmedmean.py
index 9abdbda..41648b5 100644
--- a/src/blades/aggregators/trimmedmean.py
+++ b/src/blades/aggregators/trimmedmean.py
@@ -1,30 +1,37 @@
-from typing import Union, List
+from typing import List, Union
 
 import torch
 
-from blades.client import BladesClient
+from blades.core.client import BladesClient
+
 from .mean import _BaseAggregator
 
 
 class Trimmedmean(_BaseAggregator):
-    r"""
-    A robust aggregator from paper `Byzantine-robust distributed learning: Towards optimal statistical rates <https://proceedings.mlr.press/v80/yin18a>`_
+    r"""A robust aggregator from paper `Byzantine-robust distributed learning:
+    Towards optimal statistical rates.
+
+    <https://proceedings.mlr.press/v80/yin18a>`_
+
+    It computes the coordinate-wise trimmed average of the global_model updates,
+    which can be expressed by:
 
-    It computes the coordinate-wise trimmed average of the model updates, which can be expressed by:
-    
     .. math::
              trmean := TrimmedMean( \{ \Delta_k : k \in [K] \} ),
-             
-    where the :math:`i`-th coordinate :math:`trmean_i = \frac{1}{(1-2\beta)m} \sum_{x \in U_k}x`, and :math:`U_k`
-    is a subset obtained by removing the largest and smallest :math:`\beta` fraction of its elements.
 
+    where the :math:`i`-th coordinate :math:`trmean_i = \frac{1}{(1-2\beta)m}
+    \sum_{x \in U_k}x`, and :math:`U_k`
+    is a subset obtained by removing the largest and smallest :math:`\beta` fraction
+    of its elements.
     """
-    
+
     def __init__(self, num_excluded=5):
         self.b = num_excluded
         super(Trimmedmean, self).__init__()
-    
-    def __call__(self, inputs: Union[List[BladesClient], List[torch.Tensor], torch.Tensor]):
+
+    def __call__(
+        self, inputs: Union[List[BladesClient], List[torch.Tensor], torch.Tensor]
+    ):
         updates = self._get_updates(inputs)
         if len(updates) - 2 * self.b > 0:
             b = self.b
@@ -34,12 +41,12 @@ def __call__(self, inputs: Union[List[BladesClient], List[torch.Tensor], torch.T
                 b -= 1
             if b < 0:
                 raise RuntimeError
-        
+
         largest, _ = torch.topk(updates, b, 0)
         neg_smallest, _ = torch.topk(-updates, b, 0)
         new_stacked = torch.cat([updates, -largest, neg_smallest]).sum(0)
         new_stacked /= len(updates) - 2 * b
         return new_stacked
-    
+
     def __str__(self):
         return "Trimmed Mean (b={})".format(self.b)
diff --git a/src/blades/attackers/alieclient.py b/src/blades/attackers/alieclient.py
index 2ee2de9..1cd024e 100644
--- a/src/blades/attackers/alieclient.py
+++ b/src/blades/attackers/alieclient.py
@@ -2,33 +2,43 @@
 import torch
 from scipy.stats import norm
 
-from blades.client import ByzantineClient
+from blades.core.client import ByzantineClient
 
 
 class AlieClient(ByzantineClient):
     """
-        :param num_clients: Total number of input
-        :param num_byzantine: Number of Byzantine input
+    :param num_clients: Total number of input
+    :param num_byzantine: Number of Byzantine input
     """
-    
-    def __init__(self, num_clients: int, num_byzantine: int, z=None, *args, **kwargs):
+
+    def __init__(
+        self,
+        num_clients: int,
+        num_byzantine: int,
+        z=None,
+        *args,
+        **kwargs,
+    ):
         super().__init__(*args, **kwargs)
         # Number of supporters
         if z is not None:
             self.z_max = z
         else:
             s = np.floor(num_clients / 2 + 1) - num_byzantine
-            cdf_value = (num_clients - num_byzantine - s) / (num_clients - num_byzantine)
+            cdf_value = (num_clients - num_byzantine - s) / (
+                num_clients - num_byzantine
+            )
             self.z_max = norm.ppf(cdf_value)
         self.n_good = num_clients - num_byzantine
-    
+
     def omniscient_callback(self, simulator):
-        pass 
-    
-    def local_training(self, data_batches):
         pass
 
-class AlieAdversary():
+    def local_training(self, data_batches, opt):
+        pass
+
+
+class AlieAdversary:
     def __init__(self):
         pass
 
@@ -37,7 +47,7 @@ def omniscient_callback(self, simulator):
         for client in simulator._clients.values():
             if not client.is_byzantine():
                 updates.append(client.get_update())
-        
+
         stacked_updates = torch.stack(updates, 1)
         mu = torch.mean(stacked_updates, 1)
         std = torch.std(stacked_updates, 1)
@@ -47,4 +57,3 @@ def omniscient_callback(self, simulator):
                 if update is None:
                     update = mu - std * client.z_max
                 client.save_update(update)
-
diff --git a/src/blades/attackers/attackclippedclusteringclient.py b/src/blades/attackers/attackclippedclusteringclient.py
index 26e065b..dab9eb8 100644
--- a/src/blades/attackers/attackclippedclusteringclient.py
+++ b/src/blades/attackers/attackclippedclusteringclient.py
@@ -5,22 +5,18 @@
 from sklearn.cluster import AgglomerativeClustering
 
 from blades.aggregators.clippedclustering import Clippedclustering
-from blades.client import ByzantineClient
+from blades.core.client import ByzantineClient
 
 
 class AttackclippedclusteringClient(ByzantineClient):
-    r"""
-    """
-    
+    r""""""
+
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        
-        
-        
-        
 
-class AttackclippedclusteringAdversary():
-    def __init__(self, linkage='single'):
+
+class AttackclippedclusteringAdversary:
+    def __init__(self, linkage="single"):
         self.linkage = linkage
         self.agg = Clippedclustering()
 
@@ -31,9 +27,9 @@ def chain_attack(self, simulator):
                 benign_update.append(w.get_update())
         benign_update = torch.stack(benign_update, 0)
         benign_mean = benign_update.mean(dim=0).cpu().detach().numpy()
-    
+
         np_models = benign_update.cpu().detach().numpy()
-    
+
         num = len(benign_update)
         dis_max = np.zeros((num, num))
         for i in range(num):
@@ -41,13 +37,17 @@ def chain_attack(self, simulator):
                 if i == j:
                     dis_max[i, j] = 0
                 else:
-                    dis_max[i, j] = spatial.distance.cosine(np_models[i, :], np_models[j, :])
+                    dis_max[i, j] = spatial.distance.cosine(
+                        np_models[i, :], np_models[j, :]
+                    )
         dis_max[dis_max == -inf] = 0
         dis_max[dis_max == inf] = 2
         dis_max[np.isnan(dis_max)] = 2
-        clustering = AgglomerativeClustering(affinity='precomputed', linkage='single', n_clusters=2)
+        clustering = AgglomerativeClustering(
+            affinity="precomputed", linkage="single", n_clusters=2
+        )
         clustering.fit(dis_max)
-    
+
         dis_cross = inf
         for idx_i, label_i in enumerate(clustering.labels_):
             for idx_j, label_j in enumerate(clustering.labels_):
@@ -55,14 +55,30 @@ def chain_attack(self, simulator):
                     continue
                 dis = spatial.distance.cosine(np_models[idx_i, :], np_models[idx_j, :])
                 dis_cross = min(dis_cross, dis)
-    
+
         theta_cross = np.arccos(1 - dis_cross) - 0.1
 
         flag = 1 if np.sum(clustering.labels_) > num // 2 else 0
-        # values = torch.vstack(list(model for model, label in zip(updates, clustering.labels_) if label == flag)).mean( dim=0)
-        larger_group = torch.vstack(list(model / torch.norm(model) for model, label in zip(benign_update, clustering.labels_) if label == flag)).cpu().detach().numpy()
+        # values = torch.vstack(list(global_model for global_model, label
+        # in zip(updates,
+        # clustering.labels_) if label == flag)).mean( dim=0)
+        larger_group = (
+            torch.vstack(
+                list(
+                    model / torch.norm(model)
+                    for model, label in zip(benign_update, clustering.labels_)
+                    if label == flag
+                )
+            )
+            .cpu()
+            .detach()
+            .numpy()
+        )
         # theta_cross = 0
-        dis2mean = [spatial.distance.cosine(benign_update, benign_mean) for benign_update in larger_group]
+        dis2mean = [
+            spatial.distance.cosine(benign_update, benign_mean)
+            for benign_update in larger_group
+        ]
         idx_max_dis = np.argmax(dis2mean)
         theta = np.arccos(1 - dis2mean[idx_max_dis])
         mal_update = benign_update[idx_max_dis] / torch.norm(benign_update[idx_max_dis])
@@ -72,15 +88,27 @@ def chain_attack(self, simulator):
                 if theta + theta_cross >= np.pi:
                     mal_update = -10 * benign_update.mean(dim=0)
                 else:
-                    a = (np.cos(theta + theta_cross - 1e-4) - np.sin(theta + theta_cross - 1e-4) / np.tan(theta))
-                    b = (np.cos(theta_cross - 1e-4) + np.sin(theta_cross - 1e-4) / np.tan(theta))
-                    mal_update = 10 * (a * benign_update.mean(dim=0) / torch.norm(benign_update.mean(dim=0)) + b * mal_update)
+                    a = np.cos(theta + theta_cross - 1e-4) - np.sin(
+                        theta + theta_cross - 1e-4
+                    ) / np.tan(theta)
+                    b = np.cos(theta_cross - 1e-4) + np.sin(
+                        theta_cross - 1e-4
+                    ) / np.tan(theta)
+                    mal_update = 10 * (
+                        a
+                        * benign_update.mean(dim=0)
+                        / torch.norm(benign_update.mean(dim=0))
+                        + b * mal_update
+                    )
                     theta = theta + theta_cross - 1e-4
                 w.save_update(mal_update)
-                dis = spatial.distance.cosine(mal_update.cpu().detach().numpy(), benign_update.mean(dim=0).cpu().detach().numpy())
+                dis = spatial.distance.cosine(
+                    mal_update.cpu().detach().numpy(),
+                    benign_update.mean(dim=0).cpu().detach().numpy(),
+                )
                 print(dis)
         return
-    
+
     def attack_average(self, simulator):
         benign_update = []
         for w in simulator.get_clients():
@@ -88,9 +116,9 @@ def attack_average(self, simulator):
                 benign_update.append(w.get_update())
         benign_update = torch.stack(benign_update, 0)
         benign_mean = benign_update.mean(dim=0).cpu().detach().numpy()
-    
+
         np_models = benign_update.cpu().detach().numpy()
-    
+
         num = len(benign_update)
         dis_max = np.zeros((num, num))
         for i in range(num):
@@ -98,19 +126,32 @@ def attack_average(self, simulator):
                 if i == j:
                     dis_max[i, j] = 0
                 else:
-                    dis_max[i, j] = spatial.distance.cosine(np_models[i, :], np_models[j, :])
+                    dis_max[i, j] = spatial.distance.cosine(
+                        np_models[i, :], np_models[j, :]
+                    )
         dis_max[dis_max == -inf] = 0
         dis_max[dis_max == inf] = 2
         dis_max[np.isnan(dis_max)] = 2
-        clustering = AgglomerativeClustering(affinity='precomputed', linkage='average', n_clusters=2)
+        clustering = AgglomerativeClustering(
+            affinity="precomputed",
+            linkage="average",
+            n_clusters=2,
+        )
         clustering.fit(dis_max)
-    
+
         flag = 1 if np.sum(clustering.labels_) > num // 2 else 0
-        # values = torch.vstack(list(model for model, label in zip(updates, clustering.labels_) if label == flag)).mean( dim=0)
+        # values = torch.vstack(list(global_model for global_model, label
+        # in zip(updates,
+        # clustering.labels_) if label == flag)).mean( dim=0)
         larger_group = torch.vstack(
-            list(model / torch.norm(model) for model, label in zip(benign_update, clustering.labels_) if label == flag))
+            list(
+                model / torch.norm(model)
+                for model, label in zip(benign_update, clustering.labels_)
+                if label == flag
+            )
+        )
         center_point = torch.mean(larger_group, dim=0)
-    
+
         dis_cross = 0
         dis_size = 0
         for i in range(len(clustering.labels_)):
@@ -119,9 +160,9 @@ def attack_average(self, simulator):
                     dis = spatial.distance.cosine(np_models[i, :], np_models[j, :])
                     dis_cross += dis
                     dis_size += 1
-    
+
         dis_avg = dis_cross / dis_size
-        print('avg dis', dis_avg)
+        print("avg dis", dis_avg)
         theta_cross = np.arccos(1 - dis_avg) - 0.1
         theta = np.arccos(1 - spatial.distance.cosine(center_point, benign_mean))
         print(dis_cross)
@@ -130,19 +171,28 @@ def attack_average(self, simulator):
                 if theta + theta_cross >= np.pi:
                     mal_update = -1000 * benign_update.mean(dim=0)
                 else:
-                    a = (np.cos(theta + theta_cross) - np.sin(theta + theta_cross) / np.tan(theta))
-                    b = (np.cos(theta_cross) + np.sin(theta_cross) / np.tan(theta))
-                    mal_update = 1000 * (a * benign_update.mean(dim=0) / torch.norm(benign_update.mean(dim=0)) + b * center_point / torch.norm(center_point))
+                    a = np.cos(theta + theta_cross) - np.sin(
+                        theta + theta_cross
+                    ) / np.tan(theta)
+                    b = np.cos(theta_cross) + np.sin(theta_cross) / np.tan(theta)
+                    mal_update = 1000 * (
+                        a
+                        * benign_update.mean(dim=0)
+                        / torch.norm(benign_update.mean(dim=0))
+                        + b * center_point / torch.norm(center_point)
+                    )
                 w.save_update(mal_update)
-                dis = spatial.distance.cosine(mal_update.cpu().detach().numpy(),
-                                              center_point.cpu().detach().numpy())
+                dis = spatial.distance.cosine(
+                    mal_update.cpu().detach().numpy(),
+                    center_point.cpu().detach().numpy(),
+                )
                 print(np.arccos(1 - dis))
         return
-    
+
     def omniscient_callback(self, simulator):
         if self.linkage == "single":
             return self.chain_attack(simulator)
         elif self.linkage == "average":
             return self.attack_average(simulator)
         else:
-            raise NotImplementedError(f"linkage {self.linkage} is not implemented.")
+            raise NotImplementedError(f"linkage {self.linkage} " f"is not implemented.")
diff --git a/src/blades/attackers/distancemaximizationclient.py b/src/blades/attackers/distancemaximizationclient.py
index 59ecdb3..e111b84 100644
--- a/src/blades/attackers/distancemaximizationclient.py
+++ b/src/blades/attackers/distancemaximizationclient.py
@@ -1,66 +1,78 @@
 import torch
 
-from blades.client import ByzantineClient
+from blades.core.client import ByzantineClient
+
 
 class DistancemaximizationClient(ByzantineClient):
     def omniscient_callback(self, simulator):
         pass
-    
-    def local_training(self, data_batches):
+
+    def local_training(self, data_batches, opt):
         pass
-    
-class DistancemaximizationAdversary():
-    r""" 
-    """
-    
-    def __init__(self, num_byzantine: int, agg: str, dev_type='sign', threshold=5.0, threshold_diff=1e-5, *args, **kwargs):
+
+
+class DistancemaximizationAdversary:
+    r""""""
+
+    def __init__(
+        self,
+        num_byzantine: int,
+        agg: str,
+        dev_type="sign",
+        threshold=5.0,
+        threshold_diff=1e-5,
+        *args,
+        **kwargs,
+    ):
         super().__init__(*args, **kwargs)
         self.dev_type = dev_type
         self.agg = agg
         self.threshold = threshold
         self.threshold_diff = threshold_diff
         self.num_byzantine = num_byzantine
-    
+
     def attack(self, simulator):
-        all_updates = torch.stack(list(map(lambda w: w.get_update(), simulator._clients.values())))
+        all_updates = torch.stack(
+            list(map(lambda w: w.get_update(), simulator._clients.values()))
+        )
         model_re = torch.mean(all_updates, 0)
-    
-        if self.dev_type == 'sign':
+
+        if self.dev_type == "sign":
             deviation = torch.sign(model_re)
-        elif self.dev_type == 'unit_vec':
+        elif self.dev_type == "unit_vec":
             deviation = model_re / torch.norm(model_re)
-        elif self.dev_type == 'std':
+        elif self.dev_type == "std":
             deviation = torch.std(all_updates, 0)
-    
-        lamda = torch.Tensor([self.threshold]).to(all_updates.device)
-    
+
+        lambda_ = torch.Tensor([self.threshold]).to(all_updates.device)
+
         threshold_diff = self.threshold_diff
         prev_loss = -1
-        lamda_fail = lamda
+        lamda_fail = lambda_
         lamda_succ = 0
-    
-        while torch.abs(lamda_succ - lamda) > threshold_diff:
-            mal_update = (model_re - lamda * deviation)
+
+        while torch.abs(lamda_succ - lambda_) > threshold_diff:
+            mal_update = model_re - lambda_ * deviation
             mal_updates = torch.stack([mal_update] * self.num_byzantine)
             mal_updates = torch.cat((mal_updates, all_updates), 0)
-        
+
             agg_updates = simulator.server.aggregator(mal_updates)
-        
+
             loss = torch.norm(agg_updates - model_re)
-        
+
             if prev_loss < loss:
-                lamda_succ = lamda
-                lamda = lamda + lamda_fail / 2
+                lamda_succ = lambda_
+                lambda_ = lambda_ + lamda_fail / 2
             else:
-                lamda = lamda - lamda_fail / 2
-        
+                lambda_ = lambda_ - lamda_fail / 2
+
             lamda_fail = lamda_fail / 2
             prev_loss = loss
-    
-        mal_update = (model_re - lamda_succ * deviation)
+
+        mal_update = model_re - lamda_succ * deviation
         for i, client in enumerate(simulator._clients.values()):
             if client.is_byzantine():
                 client.save_update(mal_update)
-    
+
     def omniscient_callback(self, simulator):
         self.attack(simulator)
diff --git a/src/blades/attackers/fangattackclient.py b/src/blades/attackers/fangattackclient.py
index ff059aa..a8490fb 100644
--- a/src/blades/attackers/fangattackclient.py
+++ b/src/blades/attackers/fangattackclient.py
@@ -1,26 +1,34 @@
 import numpy as np
 import torch
 
-from blades.client import ByzantineClient
 from blades.aggregators.multikrum import Multikrum
+from blades.core.client import ByzantineClient
+
 
 class FangattackClient(ByzantineClient):
     def omniscient_callback(self, simulator):
         pass
-    
-    def local_training(self, data_batches):
+
+    def local_training(self, data_batches, opt):
         pass
-    
-class FangattackAdversary():
-    r""" 
-    """
-    
-    def __init__(self, num_byzantine: int, agg: str, dev_type='std', *args, **kwargs):
+
+
+class FangattackAdversary:
+    r""""""
+
+    def __init__(
+        self,
+        num_byzantine: int,
+        agg: str,
+        dev_type="std",
+        *args,
+        **kwargs,
+    ):
         super().__init__(*args, **kwargs)
         self.dev_type = dev_type
         self.agg = agg
         self.num_byzantine = num_byzantine
-    
+
     def attack_median_and_trimmedmean(self, simulator):
         benign_update = []
         for w in simulator.get_clients():
@@ -29,39 +37,60 @@ def attack_median_and_trimmedmean(self, simulator):
         benign_update = torch.stack(benign_update, 0)
         agg_grads = torch.mean(benign_update, 0)
         deviation = torch.sign(agg_grads)
+        device = benign_update.device
         b = 2
         max_vector = torch.max(benign_update, 0)[0]
         min_vector = torch.min(benign_update, 0)[0]
-        
-        max_ = (max_vector > 0).type(torch.FloatTensor).to(benign_update.device)
-        min_ = (min_vector < 0).type(torch.FloatTensor).to(benign_update.device)
-        
+
+        max_ = (max_vector > 0).type(torch.FloatTensor).to(device)
+        min_ = (min_vector < 0).type(torch.FloatTensor).to(device)
+
         max_[max_ == 1] = b
         max_[max_ == 0] = 1 / b
         min_[min_ == 1] = b
         min_[min_ == 0] = 1 / b
-        
-        max_range = torch.cat((max_vector[:, None], (max_vector * max_)[:, None]), dim=1)
-        min_range = torch.cat(((min_vector * min_)[:, None], min_vector[:, None]), dim=1)
-        
-        rand = torch.from_numpy(np.random.uniform(0, 1, [len(deviation), self.num_byzantine])).type(torch.FloatTensor).to(benign_update.device)
-        
-        max_rand = torch.stack([max_range[:, 0]] * rand.shape[1]).T + rand * torch.stack(
-            [max_range[:, 1] - max_range[:, 0]] * rand.shape[1]).T
-        min_rand = torch.stack([min_range[:, 0]] * rand.shape[1]).T + rand * torch.stack(
-            [min_range[:, 1] - min_range[:, 0]] * rand.shape[1]).T
-        
-        mal_vec = (torch.stack(
-            [(deviation < 0).type(torch.FloatTensor)] * max_rand.shape[1]).T.to(benign_update.device) * max_rand + torch.stack(
-            [(deviation > 0).type(torch.FloatTensor)] * min_rand.shape[1]).T.to(benign_update.device) * min_rand).T
-        
+
+        max_range = torch.cat(
+            (max_vector[:, None], (max_vector * max_)[:, None]), dim=1
+        )
+        min_range = torch.cat(
+            ((min_vector * min_)[:, None], min_vector[:, None]), dim=1
+        )
+
+        rand = (
+            torch.from_numpy(
+                np.random.uniform(0, 1, [len(deviation), self.num_byzantine])
+            )
+            .type(torch.FloatTensor)
+            .to(benign_update.device)
+        )
+
+        max_rand = (
+            torch.stack([max_range[:, 0]] * rand.shape[1]).T
+            + rand * torch.stack([max_range[:, 1] - max_range[:, 0]] * rand.shape[1]).T
+        )
+        min_rand = (
+            torch.stack([min_range[:, 0]] * rand.shape[1]).T
+            + rand * torch.stack([min_range[:, 1] - min_range[:, 0]] * rand.shape[1]).T
+        )
+
+        mal_vec = (
+            torch.stack(
+                [(deviation < 0).type(torch.FloatTensor)] * max_rand.shape[1]
+            ).T.to(device)
+            * max_rand
+            + torch.stack(
+                [(deviation > 0).type(torch.FloatTensor)] * min_rand.shape[1]
+            ).T.to(device)
+            * min_rand
+        ).T
+
         for i, client in enumerate(simulator._clients.values()):
             if client.is_byzantine():
                 client.save_update(mal_vec[i])
 
     def attack_multikrum(self, simulator):
-        num_clients = len(simulator.get_clients())
-        multi_krum = Multikrum(num_clients=num_clients, num_byzantine=self.num_byzantine, k=1)
+        multi_krum = Multikrum(num_excluded=self.num_byzantine, k=1)
         benign_update = []
         for w in simulator.get_clients():
             if not w.is_byzantine():
@@ -69,52 +98,66 @@ def attack_multikrum(self, simulator):
         benign_update = torch.stack(benign_update, 0)
         agg_updates = torch.mean(benign_update, 0)
         deviation = torch.sign(agg_updates)
-        
+
         def compute_lambda(all_updates, model_re, n_attackers):
-        
+
             distances = []
             n_benign, d = all_updates.shape
             for update in all_updates:
                 distance = torch.norm((all_updates - update), dim=1)
-                distances = distance[None, :] if not len(distances) else torch.cat((distances, distance[None, :]), 0)
-        
+                distances = (
+                    distance[None, :]
+                    if not len(distances)
+                    else torch.cat((distances, distance[None, :]), 0)
+                )
+
             distances[distances == 0] = 10000
             distances = torch.sort(distances, dim=1)[0]
-            scores = torch.sum(distances[:, :n_benign - 2 - n_attackers], dim=1)
+            scores = torch.sum(
+                distances[:, : n_benign - 2 - n_attackers],
+                dim=1,
+            )
             min_score = torch.min(scores)
-            term_1 = min_score / ((n_benign - n_attackers - 1) * torch.sqrt(torch.Tensor([d]))[0])
-            max_wre_dist = torch.max(torch.norm((all_updates - model_re), dim=1)) / (torch.sqrt(torch.Tensor([d]))[0])
-        
-            return (term_1 + max_wre_dist)
-
-        all_updates = torch.stack(list(map(lambda w: w.get_update(), simulator._clients.values())))
-        lamda = compute_lambda(all_updates, agg_updates, self.num_byzantine)
-    
+            term_1 = min_score / (
+                (n_benign - n_attackers - 1) * torch.sqrt(torch.Tensor([d]))[0]
+            )
+            max_wre_dist = torch.max(torch.norm((all_updates - model_re), dim=1)) / (
+                torch.sqrt(torch.Tensor([d]))[0]
+            )
+
+            return term_1 + max_wre_dist
+
+        all_updates = torch.stack(
+            list(map(lambda w: w.get_update(), simulator._clients.values()))
+        )
+        lambda_ = compute_lambda(all_updates, agg_updates, self.num_byzantine)
+
         threshold = 1e-5
         mal_update = []
-    
-        while lamda > threshold:
-            mal_update = (-lamda * deviation)
+
+        while lambda_ > threshold:
+            mal_update = -lambda_ * deviation
             mal_updates = torch.stack([mal_update] * self.num_byzantine)
             mal_updates = torch.cat((mal_updates, all_updates), 0)
-        
+
             # print(mal_updates.shape, n_attackers)
             agg_grads, krum_candidate = multi_krum(mal_updates)
             if krum_candidate < self.num_byzantine:
-                # print('successful lamda is ', lamda)
                 return mal_update
             else:
                 mal_update = []
-        
-            lamda *= 0.5
-    
+
+            lambda_ *= 0.5
+
         if not len(mal_update):
-            mal_update = (agg_updates - lamda * deviation)
-    
+            mal_update = agg_updates - lambda_ * deviation
+
         return mal_update
-    
+
     def omniscient_callback(self, simulator):
         if self.agg in ["median", "trimmedmean"]:
             self.attack_median_and_trimmedmean(simulator)
         else:
-            raise NotImplementedError(f"Adaptive attacks to {self.agg} is not supported yet.")
+            raise NotImplementedError(
+                f"Adaptive attacks to {self.agg} " f"is not supported yet."
+            )
diff --git a/src/blades/attackers/ipmclient.py b/src/blades/attackers/ipmclient.py
index b74c004..2aad42f 100644
--- a/src/blades/attackers/ipmclient.py
+++ b/src/blades/attackers/ipmclient.py
@@ -1,16 +1,16 @@
-from blades.client import ByzantineClient
+from blades.core.client import ByzantineClient
 
 
 class IpmClient(ByzantineClient):
     def __init__(self, epsilon: float = 0.5, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.epsilon = epsilon
-        
+
     def omniscient_callback(self, simulator):
         # Loop over good workers and accumulate their updates
         updates = []
         for w in simulator.get_clients():
             if not w.is_byzantine():
                 updates.append(w.get_update())
-        
-        self.save_update(-self.epsilon * (sum(updates)) / len(updates))
\ No newline at end of file
+
+        self.save_update(-self.epsilon * (sum(updates)) / len(updates))
diff --git a/src/blades/attackers/labelflippingclient.py b/src/blades/attackers/labelflippingclient.py
index 5c1e624..80b549e 100644
--- a/src/blades/attackers/labelflippingclient.py
+++ b/src/blades/attackers/labelflippingclient.py
@@ -1,12 +1,4 @@
-import inspect
-import os
-import sys
-
-currentdir = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe())))
-parentdir = os.path.dirname(currentdir)
-sys.path.insert(0, parentdir)
-
-from blades.client import ByzantineClient
+from blades.core.client import ByzantineClient
 
 
 class LabelflippingClient(ByzantineClient):
@@ -18,9 +10,6 @@ def __init__(self, num_classes=10, *args, **kwargs):
         """
         super().__init__(*args, **kwargs)
         self.num_classes = num_classes
-    
+
     def on_train_batch_begin(self, data, target, logs=None):
         return data, self.num_classes - 1 - target
-    
-    def __str__(self) -> str:
-        return "LableFlippingWorker"
diff --git a/src/blades/attackers/mediantailoredclient.py b/src/blades/attackers/mediantailoredclient.py
index 1040700..230ee71 100644
--- a/src/blades/attackers/mediantailoredclient.py
+++ b/src/blades/attackers/mediantailoredclient.py
@@ -1,21 +1,22 @@
 import torch
 
-from blades.client import ByzantineClient
+from blades.core.client import ByzantineClient
 
 
 class MediantailoredClient(ByzantineClient):
     r"""Uploads random noise as local update. The noise is drawn from a
-    ``normal`` distribution.  The ``means`` and ``standard deviation`` are shared among all drawn elements.
+    ``normal`` distribution.  The ``means`` and ``standard deviation`` are
+    shared among all drawn elements.
 
     :param mean: the mean for all distributions
     :param std: the standard deviation for all distributions
     """
-    
-    def __init__(self, num_byzantine: int, dev_type='std', *args, **kwargs):
+
+    def __init__(self, num_byzantine: int, dev_type="std", *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.dev_type = dev_type
         self.num_byzantine = num_byzantine
-    
+
     def omniscient_callback(self, simulator):
         # compression = 'none'
         # q_level = 2
@@ -27,40 +28,39 @@ def omniscient_callback(self, simulator):
         benign_update = torch.stack(benign_update, 0)
         model_re = torch.mean(benign_update, 0)
 
-        if self.dev_type == 'unit_vec':
-            deviation = model_re / torch.norm(model_re)  # unit vector, dir opp to good dir
-        elif self.dev_type == 'sign':
+        if self.dev_type == "unit_vec":
+            # unit vector, dir opp to good dir
+            deviation = model_re / torch.norm(model_re)
+        elif self.dev_type == "sign":
             deviation = torch.sign(model_re)
-        elif self.dev_type == 'std':
+        elif self.dev_type == "std":
             deviation = torch.std(benign_update, 0)
 
-        lamda = torch.Tensor([10.0])  # compute_lambda_our(all_updates, model_re, n_attackers)
+        # compute_lambda_our(all_updates, model_re, n_attackers)
+        lambda_ = torch.Tensor([10.0])
 
         threshold_diff = 1e-5
         prev_loss = -1
-        lamda_fail = lamda
+        lamda_fail = lambda_
         lamda_succ = 0
-        iters = 0
-        while torch.abs(lamda_succ - lamda) > threshold_diff:
-            mal_update = (model_re - lamda * deviation)
+        while torch.abs(lamda_succ - lambda_) > threshold_diff:
+            mal_update = model_re - lambda_ * deviation
             mal_updates = torch.stack([mal_update] * self.num_byzantine)
             mal_updates = torch.cat((mal_updates, benign_update), 0)
-    
+
             agg_grads = torch.median(mal_updates, 0)[0]
-    
+
             loss = torch.norm(agg_grads - model_re)
-    
+
             if prev_loss < loss:
-                lamda_succ = lamda
-                lamda = lamda + lamda_fail / 2
+                lamda_succ = lambda_
+                lambda_ = lambda_ + lamda_fail / 2
             else:
-                lamda = lamda - lamda_fail / 2
-    
+                lambda_ = lambda_ - lamda_fail / 2
+
             lamda_fail = lamda_fail / 2
             prev_loss = loss
 
-        mal_update = (model_re - lamda_succ * deviation)
-        
-        self._state['saved_update'] = mal_update
-
+        mal_update = model_re - lamda_succ * deviation
 
+        self._state["saved_update"] = mal_update
diff --git a/src/blades/attackers/noiseclient.py b/src/blades/attackers/noiseclient.py
index 1e13821..d788b09 100644
--- a/src/blades/attackers/noiseclient.py
+++ b/src/blades/attackers/noiseclient.py
@@ -2,25 +2,31 @@
 
 import torch
 
-from blades.client import ByzantineClient
+from blades.core.client import ByzantineClient
 
 
 class NoiseClient(ByzantineClient):
     r"""Uploads random noise as local update. The noise is drawn from a
-    ``normal`` distribution.  The ``means`` and ``standard deviation`` are shared among all drawn elements.
-    
+    ``normal`` distribution.  The ``means`` and ``standard deviation`` are
+    shared among all drawn elements.
+
     :param mean: the mean for all distributions
     :param std: the standard deviation for all distributions
     """
-    
-    def __init__(self, mean: Optional[float] = 0.1, std: Optional[float] = 0.1, *args, **kwargs):
+
+    def __init__(
+        self,
+        mean: Optional[float] = 0.1,
+        std: Optional[float] = 0.1,
+        *args,
+        **kwargs,
+    ):
         super().__init__(*args, **kwargs)
         self._noise_mean = mean
         self._noise_std = std
-    
+
     def omniscient_callback(self, simulator):
-        noise = torch.normal(self._noise_mean,
-                            self._noise_std,
-                            size=super().get_update().shape
-                            ).to(simulator.device)
-        self.save_update(noise)
\ No newline at end of file
+        noise = torch.normal(
+            self._noise_mean, self._noise_std, size=super().get_update().shape
+        ).to(simulator.device)
+        self.save_update(noise)
diff --git a/src/blades/attackers/permutationclient.py b/src/blades/attackers/permutationclient.py
deleted file mode 100644
index abc12f0..0000000
--- a/src/blades/attackers/permutationclient.py
+++ /dev/null
@@ -1,15 +0,0 @@
-import torch
-
-from blades.client import ByzantineClient
-
-
-class PermutationClient(ByzantineClient):
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-    
-    def omniscient_callback(self, simulator):
-        # Loop over good workers and accumulate their updates
-        true_update = self.get_update()
-        r = torch.randperm(len(true_update))
-        malicious_update = true_update[r]
-        self.save_update(malicious_update)
\ No newline at end of file
diff --git a/src/blades/attackers/signflippingclient.py b/src/blades/attackers/signflippingclient.py
index f096fcd..cde349e 100644
--- a/src/blades/attackers/signflippingclient.py
+++ b/src/blades/attackers/signflippingclient.py
@@ -1,21 +1,21 @@
 import torch
 
-from blades.client import ByzantineClient
+from blades.core.client import ByzantineClient
 
 
 class SignflippingClient(ByzantineClient):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-    
-    def local_training(self, data_batches):
+
+    def local_training(self, data_batches, opt):
         for data, target in data_batches:
             data, target = data.to(self.device), target.to(self.device)
             data, target = self.on_train_batch_begin(data=data, target=target)
-            self.optimizer.zero_grad()
-            
-            output = self.model(data)
+            opt.zero_grad()
+
+            output = self.global_model(data)
             loss = torch.clamp(self.loss_func(output, target), 0, 1e5)
             loss.backward()
-            for name, p in self.model.named_parameters():
+            for name, p in self.global_model.named_parameters():
                 p.grad.data = -p.grad.data
-            self.optimizer.step()
\ No newline at end of file
+            opt.optimizer.step()
diff --git a/src/blades/builder.py b/src/blades/builder.py
new file mode 100644
index 0000000..c304cc6
--- /dev/null
+++ b/src/blades/builder.py
@@ -0,0 +1,64 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import warnings
+
+from mmcv.cnn import MODELS as MMCV_MODELS
+from mmcv.utils import Registry
+
+MODELS = Registry("models", parent=MMCV_MODELS)
+
+BACKBONES = MODELS
+NECKS = MODELS
+ROI_EXTRACTORS = MODELS
+SHARED_HEADS = MODELS
+HEADS = MODELS
+LOSSES = MODELS
+DETECTORS = MODELS
+
+
+def build_backbone(cfg):
+    """Build backbone."""
+    return BACKBONES.build(cfg)
+
+
+def build_neck(cfg):
+    """Build neck."""
+    return NECKS.build(cfg)
+
+
+def build_roi_extractor(cfg):
+    """Build roi extractor."""
+    return ROI_EXTRACTORS.build(cfg)
+
+
+def build_shared_head(cfg):
+    """Build shared head."""
+    return SHARED_HEADS.build(cfg)
+
+
+def build_head(cfg):
+    """Build head."""
+    return HEADS.build(cfg)
+
+
+def build_loss(cfg):
+    """Build loss."""
+    return LOSSES.build(cfg)
+
+
+def build_detector(cfg, train_cfg=None, test_cfg=None):
+    """Build detector."""
+    if train_cfg is not None or test_cfg is not None:
+        warnings.warn(
+            "train_cfg and test_cfg is deprecated, "
+            "please specify them in global_model",
+            UserWarning,
+        )
+    assert (
+        cfg.get("train_cfg") is None or train_cfg is None
+    ), "train_cfg specified in both outer field and global_model field "
+    assert (
+        cfg.get("test_cfg") is None or test_cfg is None
+    ), "test_cfg specified in both outer field and global_model field "
+    return DETECTORS.build(
+        cfg, default_args=dict(train_cfg=train_cfg, test_cfg=test_cfg)
+    )
diff --git a/src/blades/models/cifar10/cctnets/utils/__init__.py b/src/blades/core/__init__.py
similarity index 100%
rename from src/blades/models/cifar10/cctnets/utils/__init__.py
rename to src/blades/core/__init__.py
diff --git a/src/blades/core/actor.py b/src/blades/core/actor.py
new file mode 100644
index 0000000..a01a680
--- /dev/null
+++ b/src/blades/core/actor.py
@@ -0,0 +1,106 @@
+import ray
+
+from blades.datasets.dataset import FLDataset
+import torch
+import torch.nn as nn
+import copy
+
+
+@ray.remote
+class _RayActor(object):
+    """Ray Actor."""
+
+    def __init__(self, dataset: object, *args, **kwargs):
+        """
+        Args:
+            aggregator (callable): A callable which takes a list of tensors and
+            returns an aggregated tensor.
+            log_interval (int): Control the frequency of logging training
+             batches
+            metrics (dict): dict of metric names and their functions
+            use_cuda (bool): Use cuda or not
+            debug (bool):
+        """
+        traindls, testdls = dataset.get_dls()
+        self.dataset = FLDataset(traindls, testdls)
+
+    def set_global_model(
+        self, model: nn.Module, opt: type(torch.optim.Optimizer), lr: float
+    ) -> None:
+        r"""Deep copy the given global_model to the client.
+
+        Args:
+            model: a Torch global_model for current client.
+            opt: client optimizer
+            lr:  local learning rate
+        """
+        self.model = copy.deepcopy(model)
+        self.optimizer = opt(self.model.parameters(), lr=lr)
+
+    def set_lr(self, lr: float) -> None:
+        r"""change the learning rate of the client optimizer.
+
+        Args:
+            lr (float): target learning rate.
+        """
+
+        for g in self.optimizer.param_groups:
+            g["lr"] = lr
+
+    def local_training(self, clients, global_model, local_round, lr, **kwargs):
+        """A proxy method that provides local training for a set of clients.
+
+        Args:
+            clients: a list of clients.
+            global_model: the global global_model from server.
+            local_round: number of local SGD.
+            lr: client learning rate.
+            **kwargs:
+
+        Returns: a list of ``clients``.
+        """
+        if "dp" in kwargs and kwargs["dp"] is True:
+            assert "clip_threshold" in kwargs and "noise_factor" in kwargs
+            dp = True
+            clip_threshold = kwargs["clip_threshold"]
+            noise_factor = kwargs["noise_factor"]
+        else:
+            dp = False
+            clip_threshold = 0
+            noise_factor = 0
+        update = []
+
+        for client in clients:
+            self.model.load_state_dict(global_model.state_dict())
+            self.set_lr(lr)
+
+            self.model.train()
+            client.set_model_ref(self.model)
+            client.on_train_round_begin(self.model)
+            data = self.dataset.get_train_data(client.id(), local_round)
+
+            client.local_training(data_batches=data, opt=self.optimizer)
+
+            client.on_train_round_end(
+                dp=dp,
+                clip_threshold=clip_threshold,
+                noise_factor=noise_factor,
+            )
+            update.append(client.get_update())
+
+        return clients
+
+    def evaluate(self, clients, global_model, round_number, batch_size, metrics):
+        update = []
+        self.model.load_state_dict(global_model.state_dict())
+        for i in range(len(clients)):
+            clients[i].set_model_ref(self.model)
+            data = self.dataset.get_all_test_data(clients[i].id())
+            result = clients[i].evaluate(
+                round_number=round_number,
+                test_set=data,
+                batch_size=batch_size,
+                metrics=metrics,
+            )
+            update.append(result)
+        return update
diff --git a/src/blades/client.py b/src/blades/core/client.py
similarity index 54%
rename from src/blades/client.py
rename to src/blades/core/client.py
index b8f34bc..d0d6222 100644
--- a/src/blades/client.py
+++ b/src/blades/core/client.py
@@ -1,4 +1,3 @@
-import copy
 import logging
 from collections import defaultdict
 from typing import Optional
@@ -11,167 +10,165 @@
 
 
 class BladesClient(object):
-
     _is_byzantine: bool = False
-
     r"""Base class for all clients.
-    
+
         .. note::
             Your honest clients should also subclass this class.
     Args:
         id (str): a unique id of the client.
-        device (str): target device if specified, all parameters will be copied to that device
+        device (str): target device if specified, all parameters will be
+                        copied to that device.
     """
-    
+
     def __init__(
-            self,
-            id: Optional[str] = None,
-            device: Optional[str] = 'cpu',
+        self,
+        id: Optional[str] = None,
+        device: Optional[torch.device] = torch.device("cpu"),
     ):
         self._state = defaultdict(dict)
         self.device = device
         self._is_trusted: bool = False
-        # self.device: str = device
-        
+
         self._json_logger = logging.getLogger("stats")
         self.debug_logger = logging.getLogger("debug")
 
         self.set_id(id)
-    
+
     def set_id(self, id: str) -> None:
-        r"""Sets the unique id of the client.
-        """
+        r"""Sets the unique id of the client."""
 
         # if not isinstance(id,str):
         #     raise TypeError(f'Client _id must be str, but got {type(id)}')
         self._id = id
-    
+
     def id(self) -> str:
         r"""Returns the unique id of the client.
-        
+
         :Example:
-        
-        >>> from blades.client import BladesClient
+
+        >>> from blades.core.client import BladesClient
         >>> client = BladesClient(id='1')
         >>> client.id()
         '1'
-
         """
         return self._id
-    
+
     def getattr(self, attr):
         return getattr(self, attr)
-    
+
     def is_byzantine(self):
-        r"""Return a boolean value specifying if the client is Byzantine"""
+        r"""Return a boolean value specifying if the client is Byzantine."""
         return self._is_byzantine
-    
+
     def is_trusted(self):
         return self._is_trusted
-    
+
     def trust(self, trusted: Optional[bool] = True) -> None:
-        r"""
-        Trusts the client as an honest participant. This property is useful
+        r"""Trusts the client as an honest participant. This property is useful
         for trust-based algorithms.
-        
+
         Args:
             trusted: Boolean; whether the client is trusted; default to True.
         """
         self._is_trusted = trusted
-    
-    def set_model(self, model: nn.Module, opt: type(torch.optim.Optimizer), lr: float) -> None:
-        r''' Deep copy the given model to the client.
-        
-            .. note::
-                To improve the scalability, this API may be removed in the future,
-                
-            :param model: a Torch model for current client.
-            :param opt: client optimizer
-            :param lr:  local learning rate
-        '''
-        self.model = copy.deepcopy(model)
-        self.optimizer = opt(self.model.parameters(), lr=lr)
-    
-    def set_lr(self, lr: float) -> None:
-        r""" change the learning rate of the client optimizer.
+
+    def set_model_ref(self, model):
+        """Copy an existing global_model reference.
 
         Args:
-            lr (float): target learning rate.
-        """
+            model: ``Torch`` global_model
 
-        for g in self.optimizer.param_groups:
-            g['lr'] = lr
+        Returns:
+        """
+        self.global_model = model
 
-    def set_loss(self, loss_func='crossentropy'):
-        if loss_func == 'crossentropy':
+    def set_loss(self, loss_func="crossentropy"):
+        if loss_func == "crossentropy":
             self.loss_func = nn.modules.loss.CrossEntropyLoss()
         else:
             raise NotImplementedError
-    
-    def set_para(self, model):
-        state_dict = model.state_dict()
-        self.model.load_state_dict(state_dict)
-    
+
     def __str__(self) -> str:
         return "BladesClient"
-    
-    def on_train_round_begin(self) -> None:
-        """Called at the beginning of each local training round in `local_training` methods.
+
+    def on_train_round_begin(self, globel_model=None) -> None:
+        """Called at the beginning of each local training round in
+        `local_training` methods.
 
         Subclasses should override for any actions to run.
 
         :param logs: Dict. Aggregated metric results up until this batch.
         """
-        self._save_para()
-        self.model = self.model.to(self.device)
-        self.model.train()
-        
-    def on_train_round_end(self, dp=False, clip_threshold=None, noise_factor=None) -> None:
-        """Called at the end of local optimization.
-        """
-        update = (self._get_para(current=True) - self._get_para(current=False))
+        self._save_para(globel_model)
+
+    def on_train_round_end(
+        self, dp=False, clip_threshold=None, noise_factor=None
+    ) -> None:
+        """Called at the end of local optimization."""
+        update = self._get_para(current=True) - self._get_para(current=False)
         if dp:
-            assert clip_threshold != None
+            assert clip_threshold is not None
             clip_tensor_norm_(update, max_norm=clip_threshold)
-    
+
             sigma = noise_factor
-            noise = torch.normal(mean=0.0, std=sigma, size=update.shape).to(update.device)
+            noise = torch.normal(
+                mean=0.0,
+                std=sigma,
+                size=update.shape,
+            ).to(update.device)
             update += noise
         self.save_update(update)
-        
-    
+        self.global_model = None
+        self._state["saved_para"].clear()
+
     def on_train_batch_begin(self, data, target, logs=None):
-        """Called at the beginning of a training batch in `local_training` methods.
+        """Called at the beginning of a training batch in `local_training`
+        methods.
 
-         Subclasses should override for any actions to run.
+        Subclasses should override for any actions to run.
 
-         :param data: input of the batch data.
-         :param target: target of the batch data.
-         :param logs: Dict. Aggregated metric results up until this batch.
-         """
+        Args:
+            data: input of the batch data.
+            target: target of the batch data.
+            logs: Dict. Aggregated metric results up until this batch.
+        """
         return data, target
 
-    def local_training(self, data_batches: list) -> None:
-        r''' Local optimizaiton of the ``client``. Byzantine input can override this method to perform adversarial attack.
+    def local_training(self, data_batches: list, opt) -> None:
+        r"""Local optimizaiton of the ``client``. Byzantine input can override
+        this method to perform adversarial attack.
 
-            Args:
-                num_rounds: Number of local optimization steps.
-                data_batches: A list of training batches for local training.
-        '''
+        Args:
+            data_batches: A list of training batches for local training.
+            opt: Optimizer.
+        """
         for data, target in data_batches:
             data, target = data.to(self.device), target.to(self.device)
             data, target = self.on_train_batch_begin(data=data, target=target)
-            self.optimizer.zero_grad()
-        
-            output = self.model(data)
-            # Clamp loss value to avoid possible 'Nan' gradient with some attack types.
+            opt.zero_grad()
+
+            output = self.global_model(data)
+            # Clamp loss value to avoid possible 'Nan' gradient with some
+            # attack types.
             loss = torch.clamp(self.loss_func(output, target), 0, 1e6)
             loss.backward()
-            self.optimizer.step()
-            
+            opt.step()
+
     def evaluate(self, round_number, test_set, batch_size, metrics):
+        """Model evaluation.
+
+        Args:
+            round_number: Current global round.
+            test_set: Data set for test.
+            batch_size: Test batch size.
+            metrics: Metrics.
+
+        Returns:
+        """
+
         dataloader = DataLoader(dataset=test_set, batch_size=batch_size)
-        self.model.eval()
+        self.global_model.eval()
         r = {
             "_meta": {"type": "client_validation"},
             "E": round_number,
@@ -180,21 +177,21 @@ def evaluate(self, round_number, test_set, batch_size, metrics):
         }
         for name in metrics:
             r[name] = 0
-        
+
         with torch.no_grad():
             for _, (data, target) in enumerate(dataloader):
                 data, target = data.to(self.device), target.to(self.device)
-                output = self.model.to(self.device)(data)
+                output = self.global_model.to(self.device)(data)
                 r["Loss"] += self.loss_func(output, target).item() * len(target)
                 r["Length"] += len(target)
-                
+
                 for name, metric in metrics.items():
                     r[name] += metric(output, target) * len(target)
-        
+
         for name in metrics:
             r[name] /= r["Length"]
         r["Loss"] /= r["Length"]
-        
+
         self._json_logger.info(r)
         self.debug_logger.info(
             f"\n=> Eval Loss={r['Loss']:.4f} "
@@ -202,63 +199,69 @@ def evaluate(self, round_number, test_set, batch_size, metrics):
             + "\n"
         )
         return r
-    
+
     def get_update(self) -> torch.Tensor:
-        r'''Returns the saved update of local optimization, represented as a vector.
-        '''
+        """Returns the saved update of local optimization, represented as a
+        vector.
+
+        Returns: a vector tensor of update parameters.
+        """
         return torch.nan_to_num(self._get_saved_update())
-    
+
     def save_update(self, update: torch.Tensor) -> None:
         r"""Sets the update of the client,.
-        
-        :param update: a vector of local update
+
+        Args:
+        update: a vector of local update
         """
-        self._state['saved_update'] = torch.clone(update).detach()
-    
+        self._state["saved_update"] = torch.clone(update).detach()
+
     def _get_saved_update(self) -> torch.Tensor:
-        return self._state['saved_update']
-    
-    def _save_para(self) -> None:
-        for name, param in self.model.named_parameters():
+        return torch.Tensor(self._state["saved_update"])
+
+    def _save_para(self, model) -> None:
+        for name, param in model.named_parameters():
             if not param.requires_grad:
                 continue
-            self._state['saved_para'][name] = torch.clone(param.data).detach().to(self.device)
-    
+            self._state["saved_para"][name] = (
+                torch.clone(param.data).detach().to(self.device)
+            )
+
     def _get_para(self, current=True) -> torch.Tensor:
         layer_parameters = []
-        
-        for name, param in self.model.named_parameters():
+
+        for name, param in self.global_model.named_parameters():
             if not param.requires_grad:
                 continue
             if current:
                 layer_parameters.append(param.data.view(-1))
             else:
-                saved_param = self._state['saved_para'][name]
+                saved_param = self._state["saved_para"][name]
                 layer_parameters.append(saved_param.data.view(-1))
-        
+
         return torch.cat(layer_parameters)
 
 
 class ByzantineClient(BladesClient):
     r"""Base class for Byzantine input.
-    
-            .. note::
-                Your Byzantine input should also subclass this class, and override ``local_training`` and
-                ``omniscient_callback`` to customize your attack.
-                
-        """
+
+    .. note::     Your Byzantine input should also subclass this class, and
+    override ``local_training`` and ``omniscient_callback`` to     customize
+    your attack.
+    """
     _is_byzantine = True
-    
+
     def __int__(self, *args, **kwargs):
         super(ByzantineClient).__init__(*args, **kwargs)
-    
+
     def omniscient_callback(self, simulator):
-        r"""A method that will be registered by the simulator and execute after each communication round.
-            It allows a Byzantine client has full knowledge of the training system, e.g., updates from all
-            input. Your Byzantine client can override this method to access information from the server
-            and other input.
-            
-            :param simulator: The running simulator.
-            :type simulator: Simulator
+        r"""A method that will be registered by the simulator and execute after
+        each communication round. It allows a Byzantine client has full
+        knowledge of the training system, e.g., updates from all input. Your
+        Byzantine client can override this method to access information from
+        the server and other input.
+
+        Args:
+            simulator: The running simulator.
         """
         pass
diff --git a/src/blades/core/server.py b/src/blades/core/server.py
new file mode 100644
index 0000000..f9c746f
--- /dev/null
+++ b/src/blades/core/server.py
@@ -0,0 +1,73 @@
+from typing import Callable
+
+import torch
+
+
+class BladesServer(object):
+    r"""Simulating the server of the federated learning system.
+
+    :ivar aggregator: a callable which takes a list of tensors and returns
+            an aggregated tensor.
+    :vartype aggregator: callable
+
+    :param  optimizer: The global optimizer, which can be any optimizer
+    from Pytorch.
+    :type optimizer: torch.optim.Optimizer
+    :param model: The global global_model
+    :type model: torch.nn.Module
+    :param aggregator: a callable which takes a list of tensors and returns
+            an aggregated tensor.
+    :type aggregator: callable
+    """
+
+    def __init__(
+        self,
+        optimizer: torch.optim.Optimizer,
+        model: torch.nn.Module,
+        aggregator: Callable[[list], torch.Tensor],
+        *args,
+        **kwargs,
+    ):
+        self.optimizer = optimizer
+        self.model = model
+        self.aggregator = aggregator
+
+    def get_opt(self) -> torch.optim.Optimizer:
+        r"""Returns the global optimizer."""
+        return self.optimizer
+
+    def zero_grad(self, set_to_none: bool = False):
+        r"""Sets the gradients of all optimized :class:`torch.Tensor` s to zero.
+        It should be called before assigning pseudo-gradient.
+
+        Args:
+            set_to_none: See `Pytorch documentation <https://pytorch.org/docs/s
+            table/generated/torch.optim.Optimizer.zero_grad.html>`_.
+        """
+        self.optimizer.zero_grad(set_to_none=set_to_none)
+
+    def get_model(self) -> torch.nn.Module:
+        r"""Returns the current global global_model."""
+        return self.model
+
+    def apply_update(self, update: torch.Tensor) -> None:
+        r"""Apply a step of global optimization.
+
+            .. note::
+                The input should be a ``Tensor``, which will be converted to
+                ``pseudo-gradient`` layer by layer.
+
+        Args:
+            update: The aggregated update.
+        """
+        self.zero_grad()
+        beg = 0
+        for group in self.optimizer.param_groups:
+            for p in group["params"]:
+                if not p.requires_grad:
+                    continue
+                end = beg + len(p.data.view(-1))
+                x = update[beg:end].reshape_as(p.data)
+                p.grad = -x.clone().detach().to(p.device)
+                beg = end
+        self.optimizer.step()
diff --git a/src/blades/simulator.py b/src/blades/core/simulator.py
similarity index 58%
rename from src/blades/simulator.py
rename to src/blades/core/simulator.py
index 3c56da1..dfc7850 100644
--- a/src/blades/simulator.py
+++ b/src/blades/core/simulator.py
@@ -1,120 +1,136 @@
 import importlib
 import logging
 from time import time
-from typing import Any, Callable, Optional, Union, List, Dict
-
+from typing import Any, Callable, Dict, List, Optional, Union
+import ray
 import numpy as np
 import torch
 from ray.util import ActorPool
 from tqdm import trange
 
-from blades.actor import _RayActor
-from blades.client import BladesClient, ByzantineClient
+from blades.core.actor import _RayActor
+from blades.core.client import BladesClient, ByzantineClient
+from blades.core.server import BladesServer
 from blades.datasets.dataset import FLDataset
-from blades.server import BladesServer
-from blades.utils.utils import reset_model_weights, set_random_seed
-from blades.utils.utils import top1_accuracy, initialize_logger
+from blades.utils.utils import (
+    initialize_logger,
+    reset_model_weights,
+    set_random_seed,
+    top1_accuracy,
+)
 
 
 class Simulator(object):
     """Synchronous and parallel training with specified aggregators.
-    
+
     :param dataset: FLDataset that consists local data of all input
-    :param aggregator: String (name of build-in aggregation scheme) or
-                       a callable which takes a list of tensors and returns an aggregated tensor.
-    :param num_byzantine: Number of Byzantine input under build-in attack.
-                          It should be ``0`` if you have custom attack strategy.
+    :param aggregator: String (name of built-in aggregation scheme) or
+                       a callable which takes a list of tensors and returns an
+                       aggregated tensor.
+    :param num_byzantine: Number of Byzantine input under built-in attack.
+                          It should be `0` if you have custom attack strategy.
     :type num_byzantine: int, optional
-    :param attack: ``None`` by default. One of the build-in attacks, i.e., ``None``, ``noise``, ``labelflipping``,
+    :param attack: ``None`` by default. One of the built-in attacks, i.e.,
+                    ``None``, ``noise``, ``labelflipping``,
                     ``signflipping``, ``alie``,
-                    ``ipm``. It should be ``None`` if you have custom attack strategy.
+                    ``ipm``.
+                    It should be ``None`` if you have custom attack strategy.
     :type attack: str
-    :param num_actors: Number of ``Ray actors`` that will be created for local training.
+    :param num_actors: Number of ``Ray actors`` that will be created for local
+                training.
     :type num_actors: int
-    :param mode: Training mode, either ``actor`` or ``trainer``. For large scale client population,
-                    ``actor mode`` is favorable
+    :param mode: Training mode, either ``actor`` or ``trainer``. For large
+            scale client population, ``actor mode`` is favorable
     :type mode: str
     :param log_path: The path of logging
     :type log_path: str
-    
     """
-    
+
     def __init__(
-            self,
-            dataset: FLDataset,
-            num_byzantine: Optional[int] = 0,
-            attack: Optional[str] = None,
-            attack_kws: Optional[Dict[str, float]] = None,
-            adversary_kws: Optional[Dict[str, float]] = None,
-            aggregator: Union[Callable[[list], torch.Tensor], str] = 'mean',
-            aggregator_kws: Optional[Dict[str, float]] = None,
-            num_actors: Optional[int] = 1,
-            gpu_per_actor: Optional[float] = 0,
-            log_path: str = "./outputs",
-            metrics: Optional[dict] = None,
-            use_cuda: Optional[bool] = False,
-            seed: Optional[int] = None,
-            **kwargs,
+        self,
+        dataset: FLDataset,
+        num_byzantine: Optional[int] = 0,
+        attack: Optional[str] = None,
+        attack_kws: Optional[Dict[str, float]] = None,
+        adversary_kws: Optional[Dict[str, float]] = None,
+        aggregator: Union[Callable[[list], torch.Tensor], str] = "mean",
+        aggregator_kws: Optional[Dict[str, float]] = None,
+        num_actors: Optional[int] = 1,
+        gpu_per_actor: Optional[float] = 0,
+        log_path: str = "./outputs",
+        metrics: Optional[dict] = None,
+        use_cuda: Optional[bool] = False,
+        seed: Optional[int] = None,
+        **kwargs,
     ):
-    
+
         if not adversary_kws:
             adversary_kws = {}
         if use_cuda or ("gpu_per_actor" in kwargs and kwargs["gpu_per_actor"] > 0.0):
             self.device = torch.device("cuda")
         else:
             self.device = torch.device("cpu")
-        
+
         if aggregator_kws is None:
             aggregator_kws = {}
         self._init_aggregator(aggregator=aggregator, aggregator_kws=aggregator_kws)
-        
+
         # Setup logger
         initialize_logger(log_path)
         self.metrics = {"top1": top1_accuracy} if metrics is None else metrics
         self.json_logger = logging.getLogger("stats")
         self.debug_logger = logging.getLogger("debug")
         self.debug_logger.info(self.__str__())
-        
+
         self.random_states = {}
         self.omniscient_callbacks = []
-        
+
         if kwargs:
-            # User passed in extra keyword arguments but isn't connecting through
-            # the simulator. Raise an error, since most likely a typo in keyword
+            # User passed in extra keyword arguments but isn't connecting
+            # through the simulator. Raise an error, since most likely a typo
+            # in keyword
             unknown = ", ".join(kwargs)
             raise RuntimeError(f"Unknown keyword argument(s): {unknown}")
-        
-        self.ray_actors = [_RayActor.options(num_gpus=gpu_per_actor).remote(dataset)
-                            for _ in range(num_actors)]
+
+        self.ray_actors = [
+            _RayActor.options(num_gpus=gpu_per_actor).remote(dataset)
+            for _ in range(num_actors)
+        ]
         self.actor_pool = ActorPool(self.ray_actors)
-        
-        
+
         if type(dataset) != FLDataset:
             traindls, testdls = dataset.get_dls()
             self.dataset = FLDataset(traindls, testdls)
-        
+
         if attack_kws is None:
             attack_kws = {}
-        self._setup_clients(attack, num_byzantine=num_byzantine, attack_kws=attack_kws)
+        self._setup_clients(
+            attack,
+            num_byzantine=num_byzantine,
+            attack_kws=attack_kws,
+        )
         self._setup_adversary(attack, adversary_kws=adversary_kws)
 
         set_random_seed(seed)
-    
+
     def _init_aggregator(self, aggregator, aggregator_kws):
         if type(aggregator) == str:
-            agg_path = importlib.import_module('blades.aggregators.%s' % aggregator)
+            agg_path = importlib.import_module("blades.aggregators.%s" % aggregator)
             agg_scheme = getattr(agg_path, aggregator.capitalize())
             self.aggregator = agg_scheme(**aggregator_kws)
         else:
             self.aggregator = aggregator
-    
+
     def _setup_adversary(self, attack: str, adversary_kws):
-        module_path = importlib.import_module('blades.attackers.%sclient' % attack)
-        adversary_cls = getattr(module_path, '%sAdversary' % attack.capitalize(), lambda: None)
+        module_path = importlib.import_module("blades.attackers.%sclient" % attack)
+        adversary_cls = getattr(
+            module_path, "%sAdversary" % attack.capitalize(), lambda: None
+        )
         self.adversary = adversary_cls(**adversary_kws) if adversary_cls else None
-        
+
     def _setup_clients(self, attack: str, num_byzantine, attack_kws):
         import importlib
+
         if attack is None:
             num_byzantine = 0
         users = self.dataset.get_clients()
@@ -122,74 +138,81 @@ def _setup_clients(self, attack: str, num_byzantine, attack_kws):
         for i, u in enumerate(users):
             # u = str(u)
             if i < num_byzantine:
-                module_path = importlib.import_module('blades.attackers.%sclient' % attack)
-                attack_scheme = getattr(module_path, '%sClient' % attack.capitalize())
+                module_path = importlib.import_module(
+                    "blades.attackers.%sclient" % attack
+                )
+                attack_scheme = getattr(module_path, "%sClient" % attack.capitalize())
                 client = attack_scheme(id=u, device=self.device, **attack_kws)
                 self._register_omniscient_callback(client.omniscient_callback)
             else:
                 client = BladesClient(id=u, device=self.device)
             self._clients[u] = client
-    
+
     def _register_omniscient_callback(self, callback):
         self.omniscient_callbacks.append(callback)
-    
+
     def get_clients(self) -> List:
-        r"""Return all clients as a list.
-        """
+        r"""Return all clients as a list."""
         return list(self._clients.values())
-    
+
     def set_trusted_clients(self, ids: List[str]) -> None:
-        r"""Set a list of input as trusted. This is usable for trusted-based algorithms
-        that assume some clients are known as not Byzantine.
-    
+        r"""Set a list of input as trusted. This is usable for trusted-based
+        algorithms that assume some clients are known as not Byzantine.
+
         :param ids: a list of client ids that are trusted
         :type ids: list
         """
         for id in ids:
             self._clients[id].trust()
-    
+
     def cache_random_state(self) -> None:
         # This function should be used for reproducibility
-        if self.device != torch.device('cpu'):
+        if self.device != torch.device("cpu"):
             self.random_states["torch_cuda"] = torch.cuda.get_rng_state()
         self.random_states["torch"] = torch.get_rng_state()
         self.random_states["numpy"] = np.random.get_state()
-    
+
     def restore_random_state(self) -> None:
         # This function should be used for reproducibility
-        if self.device != torch.device('cpu'):
+        if self.device != torch.device("cpu"):
             torch.cuda.set_rng_state(self.random_states["torch_cuda"])
         torch.set_rng_state(self.random_states["torch"])
         np.random.set_state(self.random_states["numpy"])
-    
-    def register_attackers(self, clients: List[ByzantineClient], replace_indices=None) -> None:
-        r"""Register a list of clients as attackers. Those malicious clients replace the first few clients.
-        
+
+    def register_attackers(
+        self, clients: List[ByzantineClient], replace_indices=None
+    ) -> None:
+        r"""Register a list of clients as attackers. Those malicious clients
+        replace the first few clients.
+
         Args:
-            clients:  a list of Byzantine clients that will replace some of honest ones.
-            replace_indices:  a list of indices of clients to be replaced by the Byzantine clients. The length of this
-                                list should be equal to that of ``clients`` parameter. If it remains ``None``, the first ``n`` clients
-                                will be replaced, where ``n`` is the length of ``clients``.
+            clients:  a list of Byzantine clients that will replace some of
+                        honest ones.
+            replace_indices:  a list of indices of clients to be replaced by
+                                the Byzantine clients. The length of this
+                                list should be equal to that of ``clients``
+                                parameter. If it remains ``None``, the first
+                                ``n`` clients will be replaced, where ``n`` is
+                                 the length of ``clients``.
         """
         if replace_indices:
             assert len(clients) < len(replace_indices)
         else:
             replace_indices = list(range(len(clients)))
         assert len(clients) < len(self._clients)
-        
+
         client_li = self.get_clients()
         for i in replace_indices:
             id = client_li[i].id()
             clients[i].set_id(id)
             self._clients[id] = clients[i]
             self._register_omniscient_callback(clients[i].omniscient_callback)
-    
-    def parallel_call(self, clients,
-                      f: Callable[[BladesClient], None]) -> None:
+
+    def parallel_call(self, clients, f: Callable[[BladesClient], None]) -> None:
         self.cache_random_state()
         _ = [f(worker) for worker in clients]
         self.restore_random_state()
-    
+
     def parallel_get(self, clients, f: Callable[[BladesClient], Any]) -> list:
         results = []
         for w in clients:
@@ -197,99 +220,118 @@ def parallel_get(self, clients, f: Callable[[BladesClient], Any]) -> list:
             results.append(f(w))
             self.restore_random_state()
         return results
-    
-    def train_actor(self,
-                    global_round: int,
-                    num_rounds: int,
-                    clients: List[BladesClient],
-                    lr: float,
-                    *args,
-                    **kwargs,
-                    ) -> None:
-        r"""Run local training using ``ray`` actors
-        
+
+    def train_actor(
+        self,
+        global_round: int,
+        num_rounds: int,
+        clients: List[BladesClient],
+        lr: float,
+        *args,
+        **kwargs,
+    ) -> None:
+        r"""Run local training using ``ray`` actors.
+
         Args:
             global_round (int): The current global round.
             num_rounds (int): The number of local update steps.
             clients (list): A list of clients that perform local training.
             lr (float): Learning rate for client optimizer.
         """
-        
+
         # TODO: randomly select a subset of input for local training
         self.debug_logger.info(f"Train global round {global_round}")
-        
-        t_s = time()
+
         # Allocate input to actors:
-        global_model = self.server.get_model()
+        global_model = ray.put(self.server.get_model())
         client_groups = np.array_split(self.get_clients(), len(self.ray_actors))
-        all_results = self.actor_pool.map(
-            lambda actor, clients:
-            actor.local_training.remote(
+        all_results = []
+        for clients, actor in zip(client_groups, self.ray_actors):
+            ref_clients = actor.local_training.remote(
                 clients=clients,
-                model=global_model,
+                global_model=global_model,
                 local_round=num_rounds,
                 lr=lr,
                 *args,
                 **kwargs,
-            ),
-            client_groups
-        )
-        
-       
-        updates = [update for returns in list(all_results) for update in returns]
-        for client, update in zip(clients, updates):
-            client.save_update(update)
-        
+            )
+            all_results.append(ref_clients)
+
+        clients = [
+            client for client_group in all_results for client in ray.get(client_group)
+        ]
+
+        for client in clients:
+            self._clients[client._id] = client
+
         if self.adversary:
             self.adversary.omniscient_callback(self)
-            
-        # If there are Byzantine workers, ask them to craft attackers based on the updated settings.
+
+        # If there are Byzantine workers, ask them to craft attackers based on
+        # the updated settings.
         for omniscient_attacker_callback in self.omniscient_callbacks:
             omniscient_attacker_callback(self)
-        
-        updates = self.parallel_get(clients, lambda w: w.get_update())
+
+        # updates = self.parallel_get(clients, lambda w: w.get_update())
         aggregated = self.server.aggregator(clients)
         self.server.apply_update(aggregated)
         # self.log_variance(global_round, updates)
-    
+
     def test_actor(self, global_round, batch_size):
-        """Evaluates the global model using test set.
+        """Evaluates the global global_model using test set.
+
+        Args:
+            global_round: the current global round number
+            batch_size: test batch size
 
-       :param global_round: the current global round number
-       :param batch_size: test batch size
-       """
+        Returns:
+        """
         global_model = self.server.get_model()
         client_groups = np.array_split(self.get_clients(), len(self.ray_actors))
-        
+
         all_results = self.actor_pool.map(
-            lambda actor, clients:
-            actor.evaluate.remote(
+            lambda actor, clients: actor.evaluate.remote(
                 clients=clients,
-                model=global_model,
+                global_model=global_model,
                 round_number=global_round,
                 batch_size=batch_size,
                 metrics=self.metrics,
             ),
-            client_groups
+            client_groups,
         )
-        
+
         metrics = [update for returns in list(all_results) for update in returns]
-        
+
         loss, top1 = self.log_validate(metrics)
-        self.debug_logger.info(f"Test global round {global_round}, loss: {loss}, top1: {top1}")
+        self.debug_logger.info(
+            f"Test global round {global_round}, loss: {loss}, top1: {top1}"
+        )
         return loss, top1
-    
+
     def log_variance(self, cur_round, update):
         updates = []
         for client in self._clients.values():
             if not client.is_byzantine():
                 updates.append(client.get_update())
         mean_update = torch.mean(torch.vstack(updates), dim=0)
-        var_avg = torch.mean(torch.var(torch.vstack(updates), dim=0, unbiased=False)).item()
-        norm = torch.norm(torch.var(torch.vstack(updates), dim=0, unbiased=False)).item()
+        var_avg = torch.mean(
+            torch.var(torch.vstack(updates), dim=0, unbiased=False)
+        ).item()
+        norm = torch.norm(
+            torch.var(torch.vstack(updates), dim=0, unbiased=False)
+        ).item()
         avg_norm = torch.norm(mean_update)
-        var_norm = torch.sqrt(torch.mean(torch.tensor([torch.norm(model_update - mean_update) ** 2 for model_update in updates])))
-        
+        var_norm = torch.sqrt(
+            torch.mean(
+                torch.tensor(
+                    [
+                        torch.norm(model_update - mean_update) ** 2
+                        for model_update in updates
+                    ]
+                )
+            )
+        )
+
         r = {
             "_meta": {"type": "variance"},
             "Round": cur_round,
@@ -301,24 +343,29 @@ def log_variance(self, cur_round, update):
 
         # Output to file
         self.json_logger.info(r)
-        
-    
+
     def log_validate(self, metrics):
-        top1 = np.average([metric['top1'] for metric in metrics], weights=[metric['Length'] for metric in metrics])
-        loss = np.average([metric['Loss'] for metric in metrics], weights=[metric['Length'] for metric in metrics])
+        top1 = np.average(
+            [metric["top1"] for metric in metrics],
+            weights=[metric["Length"] for metric in metrics],
+        )
+        loss = np.average(
+            [metric["Loss"] for metric in metrics],
+            weights=[metric["Length"] for metric in metrics],
+        )
         r = {
             "_meta": {"type": "test"},
-            "Round": metrics[0]['E'],
+            "Round": metrics[0]["E"],
             "top1": top1,
-            "Length": np.sum([metric['Length'] for metric in metrics]),
+            "Length": np.sum([metric["Length"] for metric in metrics]),
             "Loss": loss,
         }
         self.json_logger.info(r)
         return loss, top1
-    
+
     def log_train(self, progress, batch_idx, epoch, results):
         length = sum(res["length"] for res in results)
-        
+
         r = {
             "_meta": {"type": "train"},
             "Round": epoch,
@@ -326,57 +373,60 @@ def log_train(self, progress, batch_idx, epoch, results):
             "Length": length,
             "Loss": sum(res["loss"] * res["length"] for res in results) / length,
         }
-        
+
         for metric_name in self.metrics:
             r[metric_name] = (
-                    sum(res["metrics"][metric_name] * res["length"] for res in results)
-                    / length
+                sum(res["metrics"][metric_name] * res["length"] for res in results)
+                / length
             )
-        
+
         # Output to console
         total = len(self._clients.values()[0].data_loader.dataset)
         pct = 100 * progress / total
         self.debug_logger.info(
-            f"[Round{r['E']:2}B{r['B']:<3}| {progress:6}/{total} ({pct:3.0f}%) ] Loss: {r['Loss']:.4f} "
+            f"[Round{r['E']:2}B{r['B']:<3}| {progress:6}/{total} ({pct:3.0f}%)"
+            f" ] Loss: {r['Loss']:.4f} "
             + " ".join(name + "=" + "{:>8.4f}".format(r[name]) for name in self.metrics)
         )
         # Output to file
         self.json_logger.info(r)
-    
+
     def run(
-            self,
-            model: torch.nn.Module,
-            server_optimizer: Union[torch.optim.Optimizer, str] = 'SGD',
-            client_optimizer: Union[torch.optim.Optimizer, str] = 'SGD',
-            loss: Optional[str] = 'crossentropy',
-            global_rounds: Optional[int] = 1,
-            local_steps: Optional[int] = 1,
-            validate_interval: Optional[int] = 1,
-            test_batch_size: Optional[int] = 64,
-            server_lr: Optional[float] = 0.1,
-            client_lr: Optional[float] = 0.1,
-            server_lr_scheduler: Optional[torch.optim.lr_scheduler.MultiStepLR] = None,
-            client_lr_scheduler: Optional[torch.optim.lr_scheduler.MultiStepLR] = None,
-            dp_kws: Optional[Dict[str, float]] = None,
+        self,
+        global_model: torch.nn.Module,
+        server_optimizer: Union[torch.optim.Optimizer, str] = "SGD",
+        client_optimizer: Union[torch.optim.Optimizer, str] = "SGD",
+        loss: Optional[str] = "crossentropy",
+        global_rounds: Optional[int] = 1,
+        local_steps: Optional[int] = 1,
+        validate_interval: Optional[int] = 1,
+        test_batch_size: Optional[int] = 64,
+        server_lr: Optional[float] = 0.1,
+        client_lr: Optional[float] = 0.1,
+        server_lr_scheduler: Optional[torch.optim.lr_scheduler.MultiStepLR] = None,
+        client_lr_scheduler: Optional[torch.optim.lr_scheduler.MultiStepLR] = None,
+        dp_kws: Optional[Dict[str, float]] = None,
     ):
         """Run the adversarial training.
 
-        :param model: The global model for training.
-        :type model: torch.nn.Module
-        :param server_optimizer: Pytorch optimizer for server-side optimization.
-                                    Currently, the `str` type only supports ``SGD``
+        :param global_model: The global global_model for training.
+        :type global_model: torch.nn.Module
+        :param server_optimizer: The optimizer for server-side optimization.
+                                urrently, the `str` type only supports ``SGD``
         :type server_optimizer: torch.optim.Optimizer or str
-        :param client_optimizer: Pytorch optimizer for client-side optimization.
-                                 Currently, the ``str`` type only supports ``SGD``
+        :param client_optimizer: Optimizer for client-side optimization.
+                            Currently, the ``str`` type only supports ``SGD``
         :type client_optimizer: torch.optim.Optimizer or str
-        :param loss: A Pytorch Loss function. See `Python documentation <https://pytorch.org/docs/stable/nn.html#loss-functions>`_.
+        :param loss: A Pytorch Loss function. See `Python documentation
+        <https://pytorch.org/docs/stable/nn.html#loss-functions>`_.
                      Currently, the `str` type only supports ``crossentropy``
         :type loss: str
         :param global_rounds: Number of communication rounds in total.
         :type global_rounds: int, optional
         :param local_steps: Number of local steps of each global round.
         :type local_steps: int, optional
-        :param validate_interval: Interval of evaluating the global model using test set.
+        :param validate_interval: Interval of evaluating the global global_model using
+        test set.
         :type validate_interval: int, optional
         :param test_batch_size: Batch size of evaluation
         :type test_batch_size: int, optional
@@ -385,59 +435,67 @@ def run(
         :param client_lr: Learning rate of ``client_optimizer``
         :type client_lr: float, optional
         :param server_lr_scheduler: Server learning rate scheduler
-        :type server_lr_scheduler: torch.optim.lr_scheduler.MultiStepLR, optional
+        :type server_lr_scheduler: torch.optim.lr_scheduler.MultiStepLR
         :param client_lr_scheduler: Client learning rate scheduler
-        :type client_lr_scheduler: torch.optim.lr_scheduler.MultiStepLR, optional
+        :type client_lr_scheduler: torch.optim.lr_scheduler.MultiStepLR
         :return: None
         """
         if dp_kws:
             dp_kws.update({"dp": True})
         else:
             dp_kws = {}
-        
+
         if self.device != torch.device("cpu"):
-            model = model.to("cuda")
+            global_model = global_model.to("cuda")
 
-        reset_model_weights(model)
-        if server_optimizer == 'SGD':
-            self.server_opt = torch.optim.SGD(model.parameters(), lr=server_lr, **dp_kws)
+        reset_model_weights(global_model)
+        if server_optimizer == "SGD":
+            self.server_opt = torch.optim.SGD(
+                global_model.parameters(), lr=server_lr, **dp_kws
+            )
         else:
             self.server_opt = server_optimizer
-        
+
         self.client_opt = client_optimizer
-        self.server = BladesServer(optimizer=self.server_opt,
-                                   model=model,
-                                   aggregator=self.aggregator,
-                                   )
-        
+        self.server = BladesServer(
+            optimizer=self.server_opt,
+            model=global_model,
+            aggregator=self.aggregator,
+        )
+
         self.parallel_call(self.get_clients(), lambda client: client.set_loss(loss))
         global_start = time()
         ret = []
         global_model = self.server.get_model()
-        self.parallel_call(self.get_clients(),
-                           lambda client: client.set_model(global_model, torch.optim.SGD, client_lr))
-        
-        with trange(0, global_rounds+1) as t:
+
+        for actor in self.ray_actors:
+            actor.set_global_model.remote(global_model, torch.optim.SGD, client_lr)
+
+        with trange(0, global_rounds + 1) as t:
             for global_rounds in t:
                 round_start = time()
                 if global_rounds % validate_interval == 0:
-                    loss, top1 = self.test_actor(global_round=global_rounds, batch_size=test_batch_size)
+                    loss, top1 = self.test_actor(
+                        global_round=global_rounds, batch_size=test_batch_size
+                    )
                     t.set_postfix(loss=loss, top1=top1)
-                
-                self.train_actor(global_rounds, local_steps, self.get_clients(), client_lr, **dp_kws)
+
+                self.train_actor(
+                    global_rounds, local_steps, self.get_clients(), client_lr, **dp_kws
+                )
                 if server_lr_scheduler:
                     server_lr_scheduler.step()
 
                 if client_lr_scheduler:
                     client_lr_scheduler.step()
                     client_lr = client_lr_scheduler.get_last_lr()[0]
-                
-                # client_lr = self.server_opt.param_groups[0]['lr']
-                
+
                 ret.append(time() - round_start)
-                server_lr = self.server.get_opt().param_groups[0]['lr']
+                server_lr = self.server.get_opt().param_groups[0]["lr"]
                 self.debug_logger.info(
-                    f"E={global_rounds}; Server learning rate = {server_lr:}; Client learning rate = {client_lr:}; Time cost = {time() - global_start}")
-            
-            loss, top1 = self.test_actor(global_round=global_rounds, batch_size=test_batch_size)
+                    f"E={global_rounds}; Server learning rate = {server_lr:}; "
+                    f"Client learning rate = {client_lr:}; Time cost = "
+                    f"{time() - global_start}"
+                )
+
             return ret
diff --git a/src/blades/datasets/__init__.py b/src/blades/datasets/__init__.py
index 20d7bfe..5effd7a 100644
--- a/src/blades/datasets/__init__.py
+++ b/src/blades/datasets/__init__.py
@@ -1,3 +1,6 @@
 from .basedataset import BaseDataset
 from .cifar10 import CIFAR10
+from .cifar100 import CIFAR100
 from .mnist import MNIST
+
+__all__ = ["BaseDataset", "CIFAR10", "CIFAR100", "MNIST"]
diff --git a/src/blades/datasets/basedataset.py b/src/blades/datasets/basedataset.py
index c234fa7..d648dba 100644
--- a/src/blades/datasets/basedataset.py
+++ b/src/blades/datasets/basedataset.py
@@ -8,30 +8,33 @@
 import torch
 
 from blades.utils.utils import set_random_seed
+
 from .customdataset import CustomTensorDataset
 
 logger = logging.getLogger(__name__)
 
+
 class BaseDataset(ABC):
     train_transform = None
     test_transform = None
+
     def __init__(
-            self,
-            data_root: str = './data',
-            cache_name: str = "",
-            train_bs: Optional[int] = 32,
-            iid: Optional[bool] = True,
-            alpha: Optional[float] = 0.1,
-            num_clients: Optional[int] = 20,
-            seed = 1,
+        self,
+        data_root: str = "./data",
+        cache_name: str = "",
+        train_bs: Optional[int] = 32,
+        iid: Optional[bool] = True,
+        alpha: Optional[float] = 0.1,
+        num_clients: Optional[int] = 20,
+        seed=1,
     ):
         self.train_bs = train_bs
         if cache_name == "":
             cache_name = self.__class__.__name__
-        self._data_path = os.path.join(data_root, cache_name + '.obj')
-        
-        # Meta parameters for data partitioning, for comparison with cache. Regenerate dataset
-        # if those parameters are different.
+        self._data_path = os.path.join(data_root, cache_name + ".obj")
+
+        # Meta parameters for data partitioning, for comparison with cache.
+        # Regenerate dataset if those parameters are different.
         meta_info = {
             "num_clients": num_clients,
             "data_root": data_root,
@@ -40,11 +43,11 @@ def __init__(
             "alpha": alpha,
             "seed": seed,
         }
-        
+
         regenerate = True
         if os.path.exists(self._data_path):
-            with open(self._data_path, 'rb') as f:
-                loaded_meta_info  = pickle.load(f)
+            with open(self._data_path, "rb") as f:
+                loaded_meta_info = pickle.load(f)
                 if loaded_meta_info == meta_info:
                     regenerate = False
                 else:
@@ -55,72 +58,78 @@ def __init__(
 
         if regenerate:
             returns = self.generate_datasets(data_root, iid, alpha, num_clients, seed)
-            with open(self._data_path, 'wb') as f:
+            with open(self._data_path, "wb") as f:
                 pickle.dump(meta_info, f)
                 for obj in returns:
                     pickle.dump(obj, f)
-                    
-    
+
     @abstractmethod
-    def generate_datasets(self, path='./data', iid=True, alpha=0.1, num_clients=20, seed=1):
+    def generate_datasets(
+        self, path="./data", iid=True, alpha=0.1, num_clients=20, seed=1
+    ):
         pass
-    
+
     def _preprocess_train_data(
-            self,
-            data,
-            labels,
-            batch_size,
-            seed=0
+        self, data, labels, batch_size, seed=0
     ) -> (torch.Tensor, torch.LongTensor):
         i = 0
         # The following line is needed for reproducing the randomness of transforms.
         set_random_seed(seed)
-        
+
         idx = np.random.permutation(len(labels))
         data, labels = data[idx], labels[idx]
-        
+
         while True:
             if i * batch_size >= len(labels):
                 i = 0
                 idx = np.random.permutation(len(labels))
                 data, labels = data[idx], labels[idx]
-                
+
                 continue
             else:
-                X = data[i * batch_size:(i + 1) * batch_size, :]
-                y = labels[i * batch_size:(i + 1) * batch_size]
+                X = data[i * batch_size : (i + 1) * batch_size, :]
+                y = labels[i * batch_size : (i + 1) * batch_size]
                 i += 1
                 X = torch.Tensor(X)
                 if self.train_transform:
                     X = self.train_transform(X)
                 yield X, torch.LongTensor(y)
-    
+
     def _preprocess_test_data(
-            self,
-            data,
-            labels,
+        self,
+        data,
+        labels,
     ) -> CustomTensorDataset:
         tensor_x = torch.Tensor(data)  # transform to torch tensor
         tensor_y = torch.LongTensor(labels)
-        return CustomTensorDataset(tensor_x, tensor_y, transform_list=self.train_transform)
-    
+        return CustomTensorDataset(
+            tensor_x, tensor_y, transform_list=self.train_transform
+        )
+
     # generate two lists of dataloaders for train
     def get_dls(self):
         assert os.path.isfile(self._data_path)
-        with open(self._data_path, 'rb') as f:
-            (_, train_clients, train_data, test_clients, test_data) = [pickle.load(f) for _ in range(5)]
-        
+        with open(self._data_path, "rb") as f:
+            (_, train_clients, train_data, test_clients, test_data) = [
+                pickle.load(f) for _ in range(5)
+            ]
+
         assert sorted(train_clients) == sorted(test_clients)
-        
+
         train_dls = []
         test_dls = []
         for idx, u_id in enumerate(train_clients):
-            train_dls.append(self._preprocess_train_data(data=np.array(train_data[u_id]['x']),
-                                                         labels=np.array(train_data[u_id]['y']),
-                                                         batch_size=self.train_bs,
-                                                         ))
-            test_dls.append(self._preprocess_test_data(data=np.array(test_data[u_id]['x']),
-                                                       labels=np.array(test_data[u_id]['y']),
-                                                       ))
+            train_dls.append(
+                self._preprocess_train_data(
+                    data=np.array(train_data[u_id]["x"]),
+                    labels=np.array(train_data[u_id]["y"]),
+                    batch_size=self.train_bs,
+                )
+            )
+            test_dls.append(
+                self._preprocess_test_data(
+                    data=np.array(test_data[u_id]["x"]),
+                    labels=np.array(test_data[u_id]["y"]),
+                )
+            )
         return train_dls, test_dls
-
diff --git a/src/blades/datasets/cifar10.py b/src/blades/datasets/cifar10.py
index 38410c3..9fb3446 100644
--- a/src/blades/datasets/cifar10.py
+++ b/src/blades/datasets/cifar10.py
@@ -13,7 +13,8 @@ class CIFAR10(BaseDataset):
 
     Args:
         path (string): Root directory of dataset where directory
-            ``cifar-10-batches-py`` exists or will be saved to if download is set to True.
+            ``cifar-10-batches-py`` exists or will be saved to if download is set to
+            True.
         transform (callable, optional): A function/transform that takes in an PIL image
             and returns a transformed version. E.g, ``transforms.RandomCrop``
         target_transform (callable, optional): A function/transform that takes in the
@@ -22,63 +23,72 @@ class CIFAR10(BaseDataset):
             puts it in root directory. If dataset is already downloaded, it is not
             downloaded again.
     """
+
     stats = {
         "mean": (0.4914, 0.4822, 0.4465),
         "std": (0.2023, 0.1994, 0.2010),
     }
-    test_transform = transforms.Compose([
-        transforms.Normalize(mean=stats["mean"], std=stats["std"]),
-    ])
-    
+    test_transform = transforms.Compose(
+        [
+            transforms.Normalize(mean=stats["mean"], std=stats["std"]),
+        ]
+    )
+
     img_size = 32
-    train_transform = transforms.Compose([
-        transforms.RandomResizedCrop(32, scale=(0.75, 1.0), ratio=(1.0, 1.0)),
-        transforms.RandomHorizontalFlip(p=0.5),
-        transforms.Normalize(mean=stats["mean"], std=stats["std"]),
-        transforms.RandomErasing(p=0.25)
-    ])
-    
+    train_transform = transforms.Compose(
+        [
+            transforms.RandomResizedCrop(32, scale=(0.75, 1.0), ratio=(1.0, 1.0)),
+            transforms.RandomHorizontalFlip(p=0.5),
+            transforms.Normalize(mean=stats["mean"], std=stats["std"]),
+            transforms.RandomErasing(p=0.25),
+        ]
+    )
+
     def __init__(
-            self,
-            data_root: str = './data',
-            cache_name: str = "",
-            train_bs: Optional[int] = 32,
-            iid: Optional[bool] = True,
-            alpha: Optional[float] = 0.1,
-            num_clients: Optional[int] = 20,
-            seed: Optional[int] = 1,
+        self,
+        data_root: str = "./data",
+        cache_name: str = "",
+        train_bs: Optional[int] = 32,
+        iid: Optional[bool] = True,
+        alpha: Optional[float] = 0.1,
+        num_clients: Optional[int] = 20,
+        seed: Optional[int] = 1,
+    ):
+        super(CIFAR10, self).__init__(
+            data_root, cache_name, train_bs, iid, alpha, num_clients, seed
+        )
+
+    def generate_datasets(
+        self, path="./data", iid=True, alpha=0.1, num_clients=20, seed=1
     ):
-        super(CIFAR10, self).__init__(data_root, cache_name, train_bs, iid, alpha, num_clients, seed)
-    
-    def generate_datasets(self, path='./data', iid=True, alpha=0.1, num_clients=20, seed=1):
         train_set = torchvision.datasets.CIFAR10(train=True, download=True, root=path)
         test_set = torchvision.datasets.CIFAR10(train=False, download=True, root=path)
         x_test, y_test = test_set.data, np.array(test_set.targets)
         x_train, y_train = train_set.data, np.array(train_set.targets)
-        
-        x_train = x_train.astype('float32') / 255.0
+
+        x_train = x_train.astype("float32") / 255.0
         x_train = np.transpose(x_train, (0, 3, 1, 2))
-        x_test = x_test.astype('float32') / 255.0
+        x_test = x_test.astype("float32") / 255.0
         x_test = np.transpose(x_test, (0, 3, 1, 2))
-        
+
         np.random.seed(seed)
         x_train, y_train = shuffle(x_train, y_train)
         x_test, y_test = shuffle(x_test, y_test)
-        
+
         train_user_ids = [str(id) for id in range(num_clients)]
         x_test_splits = np.split(x_test, num_clients)
         y_test_splits = np.split(y_test, num_clients)
-        
+
         if iid:
             x_train_splits = np.split(x_train, num_clients)
             y_train_splits = np.split(y_train, num_clients)
         else:
-            print('generating non-iid data')
+            print("generating non-iid data")
             min_size = 0
             K = 10
             N = y_train.shape[0]
             client_dataidx_map = {}
-            
+
             while min_size < 10:
                 proportion_list = []
                 idx_batch = [[] for _ in range(num_clients)]
@@ -86,12 +96,19 @@ def generate_datasets(self, path='./data', iid=True, alpha=0.1, num_clients=20,
                     idx_k = np.where(y_train == k)[0]
                     np.random.shuffle(idx_k)
                     proportions = np.random.dirichlet(np.repeat(alpha, num_clients))
-                    
+
                     proportions = np.array(
-                        [p * (len(idx_j) < N / num_clients) for p, idx_j in zip(proportions, idx_batch)])
+                        [
+                            p * (len(idx_j) < N / num_clients)
+                            for p, idx_j in zip(proportions, idx_batch)
+                        ]
+                    )
                     proportions = proportions / proportions.sum()
                     proportions = (np.cumsum(proportions) * len(idx_k)).astype(int)[:-1]
-                    idx_batch = [idx_j + idx.tolist() for idx_j, idx in zip(idx_batch, np.split(idx_k, proportions))]
+                    idx_batch = [
+                        idx_j + idx.tolist()
+                        for idx_j, idx in zip(idx_batch, np.split(idx_k, proportions))
+                    ]
                     min_size = min([len(idx_j) for idx_j in idx_batch])
                     proportion_list.append(proportions)
             x_train_splits, y_train_splits = [], []
@@ -100,11 +117,17 @@ def generate_datasets(self, path='./data', iid=True, alpha=0.1, num_clients=20,
                 client_dataidx_map[j] = idx_batch[j]
                 x_train_splits.append(x_train[idx_batch[j], :])
                 y_train_splits.append(y_train[idx_batch[j]])
-        
+
         test_dataset = {}
         train_dataset = {}
         for id, index in zip(train_user_ids, range(num_clients)):
-            train_dataset[id] = {'x': x_train_splits[index], 'y': y_train_splits[index].flatten()}
-            test_dataset[id] = {'x': x_test_splits[index], 'y': y_test_splits[index].flatten()}
-        
-        return train_user_ids, train_dataset, train_user_ids, test_dataset
\ No newline at end of file
+            train_dataset[id] = {
+                "x": x_train_splits[index],
+                "y": y_train_splits[index].flatten(),
+            }
+            test_dataset[id] = {
+                "x": x_test_splits[index],
+                "y": y_test_splits[index].flatten(),
+            }
+
+        return train_user_ids, train_dataset, train_user_ids, test_dataset
diff --git a/src/blades/datasets/cifar100.py b/src/blades/datasets/cifar100.py
new file mode 100644
index 0000000..157e8b8
--- /dev/null
+++ b/src/blades/datasets/cifar100.py
@@ -0,0 +1,132 @@
+from typing import Optional
+
+import numpy as np
+import torchvision
+import torchvision.transforms as transforms
+from sklearn.utils import shuffle
+
+from .basedataset import BaseDataset
+
+
+class CIFAR100(BaseDataset):
+    """`CIFAR10 <https://www.cs.toronto.edu/~kriz/cifar.html>`_ Dataset.
+
+    Args:
+        path (string): Root directory of dataset where directory
+            ``cifar-10-batches-py`` exists or will be saved to if download is `True`.
+        transform (callable, optional): A function/transform that takes in an PIL image
+            and returns a transformed version. E.g, ``transforms.RandomCrop``
+        target_transform (callable, optional): A function/transform that takes in the
+            target and transforms it.
+        download (bool, optional): If true, downloads the dataset from the internet and
+            puts it in root directory. If dataset is already downloaded, it is not
+            downloaded again.
+    """
+
+    stats = {
+        "mean": (0.5071, 0.4867, 0.4408),
+        "std": (0.2675, 0.2565, 0.2761),
+    }
+    test_transform = transforms.Compose(
+        [
+            transforms.Normalize(mean=stats["mean"], std=stats["std"]),
+        ]
+    )
+
+    img_size = 32
+    train_transform = transforms.Compose(
+        [
+            transforms.RandomResizedCrop(32, scale=(0.75, 1.0), ratio=(1.0, 1.0)),
+            transforms.RandomHorizontalFlip(p=0.5),
+            transforms.Normalize(mean=stats["mean"], std=stats["std"]),
+            transforms.RandomErasing(p=0.25),
+        ]
+    )
+
+    def __init__(
+        self,
+        data_root: str = "./data",
+        cache_name: str = "",
+        train_bs: Optional[int] = 32,
+        iid: Optional[bool] = True,
+        alpha: Optional[float] = 0.1,
+        num_clients: Optional[int] = 20,
+        seed: Optional[int] = 1,
+    ):
+        super(CIFAR100, self).__init__(
+            data_root, cache_name, train_bs, iid, alpha, num_clients, seed
+        )
+
+    def generate_datasets(
+        self, path="./data", iid=True, alpha=0.1, num_clients=20, seed=1
+    ):
+        train_set = torchvision.datasets.CIFAR10(train=True, download=True, root=path)
+        test_set = torchvision.datasets.CIFAR10(train=False, download=True, root=path)
+        x_test, y_test = test_set.data, np.array(test_set.targets)
+        x_train, y_train = train_set.data, np.array(train_set.targets)
+
+        x_train = x_train.astype("float32") / 255.0
+        x_train = np.transpose(x_train, (0, 3, 1, 2))
+        x_test = x_test.astype("float32") / 255.0
+        x_test = np.transpose(x_test, (0, 3, 1, 2))
+
+        np.random.seed(seed)
+        x_train, y_train = shuffle(x_train, y_train)
+        x_test, y_test = shuffle(x_test, y_test)
+
+        train_user_ids = [str(id) for id in range(num_clients)]
+        x_test_splits = np.split(x_test, num_clients)
+        y_test_splits = np.split(y_test, num_clients)
+
+        if iid:
+            x_train_splits = np.split(x_train, num_clients)
+            y_train_splits = np.split(y_train, num_clients)
+        else:
+            print("generating non-iid data")
+            min_size = 0
+            K = 10
+            N = y_train.shape[0]
+            client_dataidx_map = {}
+
+            while min_size < 10:
+                proportion_list = []
+                idx_batch = [[] for _ in range(num_clients)]
+                for k in range(K):
+                    idx_k = np.where(y_train == k)[0]
+                    np.random.shuffle(idx_k)
+                    proportions = np.random.dirichlet(np.repeat(alpha, num_clients))
+
+                    proportions = np.array(
+                        [
+                            p * (len(idx_j) < N / num_clients)
+                            for p, idx_j in zip(proportions, idx_batch)
+                        ]
+                    )
+                    proportions = proportions / proportions.sum()
+                    proportions = (np.cumsum(proportions) * len(idx_k)).astype(int)[:-1]
+                    idx_batch = [
+                        idx_j + idx.tolist()
+                        for idx_j, idx in zip(idx_batch, np.split(idx_k, proportions))
+                    ]
+                    min_size = min([len(idx_j) for idx_j in idx_batch])
+                    proportion_list.append(proportions)
+            x_train_splits, y_train_splits = [], []
+            for j in range(num_clients):
+                np.random.shuffle(idx_batch[j])
+                client_dataidx_map[j] = idx_batch[j]
+                x_train_splits.append(x_train[idx_batch[j], :])
+                y_train_splits.append(y_train[idx_batch[j]])
+
+        test_dataset = {}
+        train_dataset = {}
+        for id, index in zip(train_user_ids, range(num_clients)):
+            train_dataset[id] = {
+                "x": x_train_splits[index],
+                "y": y_train_splits[index].flatten(),
+            }
+            test_dataset[id] = {
+                "x": x_test_splits[index],
+                "y": y_test_splits[index].flatten(),
+            }
+
+        return train_user_ids, train_dataset, train_user_ids, test_dataset
diff --git a/src/blades/datasets/customdataset.py b/src/blades/datasets/customdataset.py
index 446c7d4..4087e38 100644
--- a/src/blades/datasets/customdataset.py
+++ b/src/blades/datasets/customdataset.py
@@ -6,16 +6,16 @@ def __init__(self, data_X, data_y, transform_list=None):
         tensors = (data_X, data_y)
         self.tensors = tensors
         self.transforms = transform_list
-    
+
     def __getitem__(self, index):
         x = self.tensors[0][index]
-        
+
         if self.transforms:
             x = self.transforms(x)
-        
+
         y = self.tensors[1][index]
-        
+
         return x, y
-    
+
     def __len__(self):
         return self.tensors[1].size(0)
diff --git a/src/blades/datasets/dataset.py b/src/blades/datasets/dataset.py
index 634606d..93fe1b3 100644
--- a/src/blades/datasets/dataset.py
+++ b/src/blades/datasets/dataset.py
@@ -1,115 +1,47 @@
 from warnings import warn
 
 
-# class BaseDataset(ABC):
-#     def __init__(
-#             self,
-#             data_root: str = './data',
-#             train_bs: Optional[int] = 32,
-#             iid: Optional[bool] = True,
-#             alpha: Optional[float] = 0.1,
-#             num_clients: Optional[int] = 20
-#     ):
-#         self.train_bs = train_bs
-#         self._data_path = os.path.join(data_root, self.__class__.__name__ + '.obj')
-#         if not os.path.exists(self._data_path):
-#             self.generate_datasets(data_root, iid, alpha, num_clients)
-#
-#     @abstractmethod
-#     def generate_datasets(self, path='./data', iid=True, alpha=0.1, num_clients=20):
-#         pass
-#
-#     @staticmethod
-#     def _preprocess_train_data(
-#             data,
-#             labels,
-#             batch_size,
-#             seed=0
-#     ) -> (torch.Tensor, torch.LongTensor):
-#         i = 0
-#         # The following line is needed for reproducing the randomness of transforms.
-#         set_random_seed(seed)
-#
-#         idx = np.random.permutation(len(labels))
-#         data, labels = data[idx], labels[idx]
-#
-#         while True:
-#             if i * batch_size >= len(labels):
-#                 i = 0
-#                 idx = np.random.permutation(len(labels))
-#                 data, labels = data[idx], labels[idx]
-#
-#                 continue
-#             else:
-#                 X = data[i * batch_size:(i + 1) * batch_size, :]
-#                 y = labels[i * batch_size:(i + 1) * batch_size]
-#                 i += 1
-#                 X = torch.Tensor(X)
-#                 yield X, torch.LongTensor(y)
-#
-#     @staticmethod
-#     def _preprocess_test_data(
-#             data,
-#             labels,
-#     ) -> CustomTensorDataset:
-#         tensor_x = torch.Tensor(data)  # transform to torch tensor
-#         tensor_y = torch.LongTensor(labels)
-#         return CustomTensorDataset(tensor_x, tensor_y)
-#
-#     # generate two lists of dataloaders for train
-#     def get_dls(self):
-#         assert os.path.isfile(self._data_path)
-#         with open(self._data_path, 'rb') as f:
-#             (train_clients, train_data, test_clients, test_data) = [pickle.load(f) for _ in range(4)]
-#
-#         assert sorted(train_clients) == sorted(test_clients)
-#
-#         train_dls = []
-#         test_dls = []
-#         for idx, u_id in enumerate(train_clients):
-#             train_dls.append(self._preprocess_train_data(data=np.array(train_data[u_id]['x']),
-#                                                          labels=np.array(train_data[u_id]['y']),
-#                                                          batch_size=self.train_bs,
-#                                                          ))
-#             test_dls.append(self._preprocess_test_data(data=np.array(test_data[u_id]['x']),
-#                                                        labels=np.array(test_data[u_id]['y']),
-#                                                        ))
-#         return train_dls, test_dls
-#
-    
 class FLDataset(object):
-    """ Federated Larning dataset
+    """Federated Larning dataset.
 
     Args:
         train_dataloaders (list): dataloaders of training data for input
-        test_dataloaders (list, optional): dataloaders of test data for input. Two lists of dataloders should be
-            in corresponding order to the client they belong to.
+        test_dataloaders (list, optional): dataloaders of test data for input.
+            Two lists of dataloders should be in corresponding order to the
+             client they belong to.
     """
-    
+
     def __init__(
-            self,
-            train_dataloaders: list,
-            test_dataloaders: list = None
-            # Input: train and test dataloaders for each client in corresponding order
+        self,
+        train_dataloaders: list,
+        test_dataloaders: list = None
+        # Input: train and test dataloaders for each client in corresponding order
     ) -> None:
         if not test_dataloaders:
-            warn("No test data is given. Model evaluation will be based on train data. ")
+            warn(
+                "No test data is given. Model evaluation will be based on train data. "
+            )
             test_dataloaders = train_dataloaders
         if len(train_dataloaders) != len(test_dataloaders):
-            raise Exception("Invalid Input: Numbers of train dataloaders and test dataloaders should be equal. ")
+            raise Exception(
+                "Invalid Input: Numbers of train dataloaders and test dataloaders "
+                "should be equal. "
+            )
         self._train_dls = {}
         self._test_dls = {}
-        for idx, (traindl, testdl) in enumerate(zip(train_dataloaders, test_dataloaders)):
+        for idx, (traindl, testdl) in enumerate(
+            zip(train_dataloaders, test_dataloaders)
+        ):
             self._train_dls[idx] = traindl
             self._test_dls[idx] = testdl
         self._clients = list(range(len(self._train_dls)))
-    
+
     def get_clients(self):
         return self._clients
-    
+
     def get_train_data(self, u_id, num_batches):
         data = [next(self._train_dls[u_id]) for _ in range(num_batches)]
         return data
-    
+
     def get_all_test_data(self, u_id):
         return self._test_dls[u_id]
diff --git a/src/blades/datasets/mnist.py b/src/blades/datasets/mnist.py
index 2b9bd3d..5ad139d 100644
--- a/src/blades/datasets/mnist.py
+++ b/src/blades/datasets/mnist.py
@@ -8,48 +8,51 @@
 
 
 class MNIST(BaseDataset):
-    """
-    """
-    
+    """"""
+
     def __init__(
-            self,
-            data_root: str = './data',
-            cache_name: str = "",
-            train_bs: Optional[int] = 32,
-            iid: Optional[bool] = True,
-            alpha: Optional[float] = 0.1,
-            num_clients: Optional[int] = 20,
-            seed: Optional[int] = 1,
+        self,
+        data_root: str = "./data",
+        cache_name: str = "",
+        train_bs: Optional[int] = 32,
+        iid: Optional[bool] = True,
+        alpha: Optional[float] = 0.1,
+        num_clients: Optional[int] = 20,
+        seed: Optional[int] = 1,
+    ):
+        super(MNIST, self).__init__(
+            data_root, cache_name, train_bs, iid, alpha, num_clients, seed
+        )
+
+    def generate_datasets(
+        self, path="./data", iid=True, alpha=0.1, num_clients=20, seed=1
     ):
-        super(MNIST, self).__init__(data_root, cache_name, train_bs, iid, alpha, num_clients, seed)
-    
-    def generate_datasets(self, path='./data', iid=True, alpha=0.1, num_clients=20, seed=1):
         train_set = torchvision.datasets.MNIST(train=True, download=True, root=path)
         test_set = torchvision.datasets.MNIST(train=False, download=True, root=path)
         x_test, y_test = test_set.data.numpy(), test_set.targets.numpy()
         x_train, y_train = train_set.data.numpy(), train_set.targets.numpy()
-        
-        x_train = x_train.astype('float32') / 255.0
-        x_test = x_test.astype('float32') / 255.0
-        
+
+        x_train = x_train.astype("float32") / 255.0
+        x_test = x_test.astype("float32") / 255.0
+
         np.random.seed(seed)
         x_train, y_train = shuffle(x_train, y_train)
         x_test, y_test = shuffle(x_test, y_test)
-        
+
         train_user_ids = [str(id) for id in range(num_clients)]
         x_test_splits = np.split(x_test, num_clients)
         y_test_splits = np.split(y_test, num_clients)
-        
+
         if iid:
             x_train_splits = np.split(x_train, num_clients)
             y_train_splits = np.split(y_train, num_clients)
         else:
-            print('generating non-iid data')
+            print("generating non-iid data")
             min_size = 0
             K = 10
             N = y_train.shape[0]
             client_dataidx_map = {}
-            
+
             while min_size < 10:
                 proportion_list = []
                 idx_batch = [[] for _ in range(num_clients)]
@@ -57,12 +60,19 @@ def generate_datasets(self, path='./data', iid=True, alpha=0.1, num_clients=20,
                     idx_k = np.where(y_train == k)[0]
                     np.random.shuffle(idx_k)
                     proportions = np.random.dirichlet(np.repeat(alpha, num_clients))
-                    
+
                     proportions = np.array(
-                        [p * (len(idx_j) < N / num_clients) for p, idx_j in zip(proportions, idx_batch)])
+                        [
+                            p * (len(idx_j) < N / num_clients)
+                            for p, idx_j in zip(proportions, idx_batch)
+                        ]
+                    )
                     proportions = proportions / proportions.sum()
                     proportions = (np.cumsum(proportions) * len(idx_k)).astype(int)[:-1]
-                    idx_batch = [idx_j + idx.tolist() for idx_j, idx in zip(idx_batch, np.split(idx_k, proportions))]
+                    idx_batch = [
+                        idx_j + idx.tolist()
+                        for idx_j, idx in zip(idx_batch, np.split(idx_k, proportions))
+                    ]
                     min_size = min([len(idx_j) for idx_j in idx_batch])
                     proportion_list.append(proportions)
             x_train_splits, y_train_splits = [], []
@@ -71,11 +81,17 @@ def generate_datasets(self, path='./data', iid=True, alpha=0.1, num_clients=20,
                 client_dataidx_map[j] = idx_batch[j]
                 x_train_splits.append(x_train[idx_batch[j], :])
                 y_train_splits.append(y_train[idx_batch[j]])
-        
+
         test_dataset = {}
         train_dataset = {}
         for id, index in zip(train_user_ids, range(num_clients)):
-            train_dataset[id] = {'x': x_train_splits[index], 'y': y_train_splits[index].flatten()}
-            test_dataset[id] = {'x': x_test_splits[index], 'y': y_test_splits[index].flatten()}
-        
-        return train_user_ids, train_dataset, train_user_ids, test_dataset
\ No newline at end of file
+            train_dataset[id] = {
+                "x": x_train_splits[index],
+                "y": y_train_splits[index].flatten(),
+            }
+            test_dataset[id] = {
+                "x": x_test_splits[index],
+                "y": y_test_splits[index].flatten(),
+            }
+
+        return train_user_ids, train_dataset, train_user_ids, test_dataset
diff --git a/src/blades/examples/Simulation on MNIST.py b/src/blades/examples/Simulation on MNIST.py
index 0daa224..13e32ad 100644
--- a/src/blades/examples/Simulation on MNIST.py
+++ b/src/blades/examples/Simulation on MNIST.py
@@ -9,36 +9,36 @@
 import ray
 import seaborn as sns
 
+from blades.core.simulator import Simulator
 from blades.datasets import MNIST
 from blades.models.mnist import MLP
-from blades.simulator import Simulator
 
 # Initialize Ray
 ray.init()
 # ray.init(address='auto')
 
-
-# mnist = MNIST(data_root="/dbfs/data", train_bs=32, num_clients=20)  # built-in federated MNIST dataset
-mnist = MNIST(data_root="./data", train_bs=32, num_clients=20, seed=0)  # built-in federated MNIST dataset
+# mnist = MNIST(data_root="/dbfs/data", train_bs=32, num_clients=20)
+# built-in federated MNIST dataset
+mnist = MNIST(data_root="./data", train_bs=32, num_clients=20, seed=0)
 
 # configuration parameters
 conf_params = {
     "dataset": mnist,
-#     "aggregator": "trimmedmean",  # aggregation
+    #     "aggregator": "trimmedmean",  # aggregation
     "num_byzantine": 8,  # number of Byzantine input
     "attack": "ipm",  # attack strategy
     # "log_path": "dbfs/outputs",
-    "attack_kws": {   
-                          "epsilon": 100,
-                     },
+    "attack_kws": {
+        "epsilon": 100,
+    },
     "num_actors": 1,  # number of training actors
     "seed": 1,  # reproducibility
 }
 
 run_params = {
-#     "model": model,  # global model
-    "server_optimizer": 'SGD',  # ,server_opt  # server optimizer
-    "client_optimizer": 'SGD',  # client optimizer
+    #     "global_model": global_model,  # global global_model
+    "server_optimizer": "SGD",  # ,server_opt  # server optimizer
+    "client_optimizer": "SGD",  # client optimizer
     "loss": "crossentropy",  # loss function
     "global_rounds": 10,  # number of global rounds
     "local_steps": 10,  # number of steps per round
@@ -47,20 +47,19 @@
 }
 
 aggs = {
-    'mean': {},
-    'trimmedmean': {"num_byzantine": 8},
-    'geomed': {},
-    'median': {},
-    'clippedclustering': {},
+    "mean": {},
+    "trimmedmean": {"num_byzantine": 8},
+    "geomed": {},
+    "median": {},
+    "clippedclustering": {},
 }
 
-
 for agg in aggs:
-    conf_params['aggregator'] = agg
-    conf_params['log_path'] = f"./outputs/{agg}"
-#     conf_params['log_path'] = f"dbfs/outputs/{k}"
+    conf_params["aggregator"] = agg
+    conf_params["log_path"] = f"./outputs/{agg}"
+    #     conf_params['log_path'] = f"dbfs/outputs/{k}"
     model = MLP()
-    run_params['model'] = model
+    run_params["global_model"] = model
     simulator = Simulator(**conf_params)
     simulator.run(**run_params)
 
@@ -69,27 +68,26 @@ def read_json(path):
     validation = []
     with open(path, "r") as f:
         for line in f:
-            line=line.strip().replace("'", '"')
+            line = line.strip().replace("'", '"')
             line = line.replace("nan", '"nan"')
             try:
                 data = json.loads(line)
-            except:
+            except IOError:
                 print(line)
-                raise
-            if data['_meta']['type'] == 'test':
+            if data["_meta"]["type"] == "test":
                 validation.append(data)
     return validation
 
-def transform(entry, agg):  
+
+def transform(entry, agg):
     return {
-        'Round Number': entry['Round'],
-        'Accuracy (%)': entry['top1'],
-        "Loss": entry['Loss'],
-        'AGG': agg,
+        "Round Number": entry["Round"],
+        "Accuracy (%)": entry["top1"],
+        "Loss": entry["Loss"],
+        "AGG": agg,
     }
 
 
-
 df = []
 for agg in aggs:
     path = f"./outputs/{agg}/stats"
@@ -97,14 +95,10 @@ def transform(entry, agg):
     df += list(map(lambda x: transform(x, agg=agg), validation_entries))
 df = pd.DataFrame(df)
 
-
 g = sns.lineplot(
-    data=df, 
-    x="Round Number", y="Accuracy (%)",  
+    data=df,
+    x="Round Number",
+    y="Accuracy (%)",
     hue="AGG",
     ci=None,
 )
-
-
-
-
diff --git a/src/blades/examples/customize_attack.py b/src/blades/examples/customize_attack.py
index 1a2fefb..b11659e 100644
--- a/src/blades/examples/customize_attack.py
+++ b/src/blades/examples/customize_attack.py
@@ -2,58 +2,63 @@
 Customization of attack strategy
 =================================
 
-To customize attack strategies, you only need to subclass ``ByzantineClient`` and override its methods.
-At present, there are three methods for the customization of attack strategies, i.e.,
+To customize attack strategies, you only need to subclass ``ByzantineClient``
+and override its methods. At present, there are three methods for the
+customization of attack strategies, i.e.,
 
 - ``local_training``:
-    You can customize the local training process and do whatever you want. For example, flipping the
+    You can customize the local training process and do whatever you want.
+    For example, flipping the
     sign of gradients at each step.
 - ``on_train_batch_begin``:
     This method is called right before each batch, making it possible to modify
     the batch data for updating.
 - ``omniscient_callback``:
-    This method is called after local optimization. By overriding it, the attacker can
-    have full knowledge of the whole system (e.g., updates from all input), so that it can adjust the model update
-    accordingly. This method is especially useful for adaptive attacks.
+ This method is called after local optimization. By overriding it,
+ the attacker can have full knowledge of the whole system (e.g., updates from
+  all input), so that it can adjust the global_model update accordingly. This method
+   is especially useful for adaptive attacks.
 """
 
-
 import ray
 import torch
 
-from blades.client import ByzantineClient
+from blades.core.client import ByzantineClient
+from blades.core.simulator import Simulator
 from blades.datasets import MNIST
 from blades.models.mnist import MLP
-from blades.simulator import Simulator
 
 # built-in federated MNIST dataset
 mnist = MNIST(data_root="./data", train_bs=32, num_clients=10)
 
+
 # Subclass the ``ByzantineClient``
 class MaliciousClient(ByzantineClient):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.num_classes = 10
-    
-    # Attack by flipping the sign of gradient, which is equivalent to stochastic gradient ascent.
-    def local_training(self, data_batches):
+
+    # Attack by flipping the sign of gradient, which is equivalent
+    # to stochastic gradient ascent.
+    def local_training(self, data_batches, opt):
         for data, target in data_batches:
             data, target = data.to(self.device), target.to(self.device)
             data, target = self.on_train_batch_begin(data=data, target=target)
-            self.optimizer.zero_grad()
-        
-            output = self.model(data)
+            opt.zero_grad()
+
+            output = self.global_model(data)
             loss = torch.clamp(self.loss_func(output, target), 0, 1e5)
             loss.backward()
-            for name, p in self.model.named_parameters():
+            for name, p in self.global_model.named_parameters():
                 p.grad.data = -p.grad.data
-            self.optimizer.step()
-            
+            opt.step()
+
     # Attack by flipping the labels of training samples.
     def on_train_batch_begin(self, data, target, logs=None):
         return data, self.num_classes - 1 - target
-    
-    # Access the updates from all honest clients and design malicious updates accordingly.
+
+    # Access the updates from all honest clients and design
+    # malicious updates accordingly.
     def omniscient_callback(self, simulator):
         updates = []
         for w in simulator.get_clients():
@@ -73,7 +78,6 @@ def omniscient_callback(self, simulator):
 ray.init(num_gpus=0, local_mode=True)
 simulator = Simulator(**conf_params)
 
-
 # %%
 # Register attacks in the simulator.
 
@@ -85,9 +89,9 @@ def omniscient_callback(self, simulator):
 # Configure run time parameters and run the experiment.
 
 run_params = {
-    "model": MLP(),  # global model
-    "server_optimizer": 'SGD',  # server optimizer
-    "client_optimizer": 'SGD',  # client optimizer
+    "global_model": MLP(),  # global global_model
+    "server_optimizer": "SGD",  # server optimizer
+    "client_optimizer": "SGD",  # client optimizer
     "loss": "crossentropy",  # loss function
     "global_rounds": 400,  # number of global rounds
     "local_steps": 50,  # number of steps per round
diff --git a/src/blades/examples/mini_example.py b/src/blades/examples/mini_example.py
index 9542509..3717e59 100644
--- a/src/blades/examples/mini_example.py
+++ b/src/blades/examples/mini_example.py
@@ -6,15 +6,15 @@
 
 import ray
 
+from blades.core.simulator import Simulator
 from blades.datasets import MNIST
 from blades.models.mnist import MLP
-from blades.simulator import Simulator
 
 # os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
 # os.environ["CUDA_VISIBLE_DEVICES"]="2,3"
 
-
-mnist = MNIST(data_root="./data", train_bs=32, num_clients=10)  # built-in federated MNIST dataset
+# built-in federated MNIST dataset
+mnist = MNIST(data_root="./data", train_bs=32, num_clients=10)
 
 # configuration parameters
 conf_params = {
@@ -22,12 +22,11 @@
     "aggregator": "mean",  # aggregation
     "num_byzantine": 4,  # number of Byzantine input
     "attack": "alie",  # attack strategy
-    "attack_kws": {"num_clients": 10,  # attacker parameters
-                     "num_byzantine": 4},
+    "attack_kws": {"num_clients": 10, "num_byzantine": 4},  # attacker parameters
     "num_actors": 4,  # number of training actors
     # "num_actors": 10,  # number of training actors
     "use_cuda": False,
-    "gpu_per_actor": 0.,
+    "gpu_per_actor": 0.0,
     "seed": 1,  # reproducibility
 }
 
@@ -37,9 +36,9 @@
 model = MLP()
 # runtime parameters
 run_params = {
-    "model": model,  # global model
-    "server_optimizer": 'SGD',  # ,server_opt  # server optimizer
-    "client_optimizer": 'SGD',  # client optimizer
+    "global_model": model,  # global global_model
+    "server_optimizer": "SGD",  # ,server_opt  # server optimizer
+    "client_optimizer": "SGD",  # client optimizer
     "loss": "crossentropy",  # loss function
     "global_rounds": 100,  # number of global rounds
     "local_steps": 50,  # number of steps per round
diff --git a/src/blades/examples/plot_comparing_aggregation_schemes.py b/src/blades/examples/plot_comparing_aggregation_schemes.py
index c5697cb..5e6fc2c 100644
--- a/src/blades/examples/plot_comparing_aggregation_schemes.py
+++ b/src/blades/examples/plot_comparing_aggregation_schemes.py
@@ -1,18 +1,27 @@
 """
-Comparing Build-in Aggregation Schemes
+Comparing Buitd-in Aggregation Schemes
 ==========================================
-This example demonstrates the comparison of eight build-in aggregation schemes.
+This example demonstrates the comparison of eight built-in aggregation schemes.
 We draw 100 samples from two normal distributions with different mean and co-variance.
-The samples are then aggregated using the build-in aggregation rules.
+The samples are then aggregated using the built-in aggregation rules.
 """
 
 import matplotlib.pyplot as plt
 import numpy as np
 import torch
 
-from blades.aggregators import Multikrum, Median, Geomed, Autogm, Mean, Clustering, Trimmedmean, Clippedclustering
+from blades.aggregators import (
+    Autogm,
+    Clippedclustering,
+    Clustering,
+    Geomed,
+    Mean,
+    Median,
+    Multikrum,
+    Trimmedmean,
+)
 
-plt.rcParams['axes.linewidth'] = 1.5 #set the value globally
+plt.rcParams["axes.linewidth"] = 1.5  # set the value globally
 plt.rcParams["font.weight"] = "bold"
 plt.rcParams["font.size"] = 16
 plt.rcParams["axes.labelweight"] = "bold"
@@ -31,28 +40,51 @@
 
 all_data_tensor = torch.Tensor(np.concatenate([benign, outliers]))
 
-
-aggs = [Mean(),
-        Multikrum(len(all_data), len(outliers), 10),
-        Geomed(),
-        Median(),
-        Autogm(lamb=1.0),
-        Trimmedmean(num_byzantine=len(outliers)),
-        Clustering(),
-        Clippedclustering(),
-       ]
+aggs = [
+    Mean(),
+    Multikrum(len(outliers), 10),
+    Geomed(),
+    Median(),
+    Autogm(lamb=1.0),
+    Trimmedmean(num_excluded=len(outliers)),
+    Clustering(),
+    Clippedclustering(),
+]
 
 # sphinx_gallery_thumbnail_number = 1
-fig, axs = plt.subplots(figsize=(8,8))
-
+fig, axs = plt.subplots(figsize=(8, 8))
 
 ax = axs
-ax.scatter(benign[:, 0], benign[:, 1], s=sample_sz, alpha=0.6, color='r', linewidths=0.2, edgecolors='black')
-ax.scatter(outliers[:, 0], outliers[:, 1], s=sample_sz, color=[0.,  0.7, 0.,  1. ], linewidths=0.2, edgecolors='black')
+ax.scatter(
+    benign[:, 0],
+    benign[:, 1],
+    s=sample_sz,
+    alpha=0.6,
+    color="r",
+    linewidths=0.2,
+    edgecolors="black",
+)
+ax.scatter(
+    outliers[:, 0],
+    outliers[:, 1],
+    s=sample_sz,
+    color=[0.0, 0.7, 0.0, 1.0],
+    linewidths=0.2,
+    edgecolors="black",
+)
+
 
 def plot_agg(ax, agg):
     target = agg(all_data_tensor).cpu().detach().numpy()
-    ax.scatter(target[0], target[1], s=sz * 3, label=type(agg).__name__, linewidths=0.3, edgecolors='black')
+    ax.scatter(
+        target[0],
+        target[1],
+        s=sz * 3,
+        label=type(agg).__name__,
+        linewidths=0.3,
+        edgecolors="black",
+    )
+
 
 list(map(lambda agg: plot_agg(ax, agg), aggs))
 ax.set_xticks([])
@@ -63,5 +95,5 @@ def plot_agg(ax, agg):
 plt.show()
 
 # %%
-# In this example, the results of Mean and Clustering deviated away by the outliers. All the other are inside
-# the range of benign data.
\ No newline at end of file
+# In this example, the results of Mean and Clustering deviated away by the
+# outliers. All the other are inside the range of benign data.
diff --git a/src/blades/examples/todo_customize_aggregator.py b/src/blades/examples/todo_customize_aggregator.py
index 521f015..9e542bf 100644
--- a/src/blades/examples/todo_customize_aggregator.py
+++ b/src/blades/examples/todo_customize_aggregator.py
@@ -8,14 +8,15 @@
 import ray
 import torch
 
+from blades.core.simulator import Simulator
 from blades.datasets import CIFAR10
 from blades.models.cifar10 import CCTNet
-from blades.simulator import Simulator
 
-cifar10 = CIFAR10(num_clients=20, iid=True)  # built-in federated cifar10 dataset
+# built-in federated cifar10 dataset
+cifar10 = CIFAR10(num_clients=20, iid=True)
 
 
-class Median():
+class Median:
     def __call__(self, inputs):
         stacked = torch.stack(inputs, dim=0)
         values_upper, _ = stacked.median(dim=0)
@@ -37,9 +38,9 @@ def __call__(self, inputs):
 
 # runtime parameters
 run_params = {
-    "model": CCTNet(),  # global model
-    "server_optimizer": 'SGD',  # server optimizer
-    "client_optimizer": 'SGD',  # client optimizer
+    "global_model": CCTNet(),  # global global_model
+    "server_optimizer": "SGD",  # server optimizer
+    "client_optimizer": "SGD",  # client optimizer
     "loss": "crossentropy",  # loss function
     "global_rounds": 400,  # number of global rounds
     "local_steps": 20,  # number of steps per round
diff --git a/src/blades/models/backbones/__init__.py b/src/blades/models/backbones/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/src/blades/models/cifar10/cctnets/__init__.py b/src/blades/models/backbones/cctnets/__init__.py
similarity index 63%
rename from src/blades/models/cifar10/cctnets/__init__.py
rename to src/blades/models/backbones/cctnets/__init__.py
index be499dd..625203a 100644
--- a/src/blades/models/cifar10/cctnets/__init__.py
+++ b/src/blades/models/backbones/cctnets/__init__.py
@@ -1,3 +1,4 @@
 from .cct import *
 from .cvt import *
-from .vit import *
+
+# from .vit import *
diff --git a/src/blades/models/backbones/cctnets/cct.py b/src/blades/models/backbones/cctnets/cct.py
new file mode 100644
index 0000000..2398cdf
--- /dev/null
+++ b/src/blades/models/backbones/cctnets/cct.py
@@ -0,0 +1,655 @@
+import torch.nn as nn
+from torch.hub import load_state_dict_from_url
+
+from .utils.helpers import fc_check, pe_check
+from .utils.tokenizer import Tokenizer
+from .utils.transformers import TransformerClassifier
+
+try:
+    from timm.models.registry import register_model
+except ImportError:
+    from .registry import register_model
+
+model_urls = {
+    "cct_7_3x1_32": "http://ix.cs.uoregon.edu/~alih/compact-transformers/checkpoints/pretrained/cct_7_3x1_32_cifar10_300epochs.pth",
+    "cct_7_3x1_32_sine": "http://ix.cs.uoregon.edu/~alih/compact-transformers/checkpoints/pretrained/cct_7_3x1_32_sine_cifar10_5000epochs.pth",
+    "cct_7_3x1_32_c100": "http://ix.cs.uoregon.edu/~alih/compact-transformers/checkpoints/pretrained/cct_7_3x1_32_cifar100_300epochs.pth",
+    "cct_7_3x1_32_sine_c100": "http://ix.cs.uoregon.edu/~alih/compact-transformers/checkpoints/pretrained/cct_7_3x1_32_sine_cifar100_5000epochs.pth",
+    "cct_7_7x2_224_sine": "http://ix.cs.uoregon.edu/~alih/compact-transformers/checkpoints/pretrained/cct_7_7x2_224_flowers102.pth",
+    "cct_14_7x2_224": "http://ix.cs.uoregon.edu/~alih/compact-transformers/checkpoints/pretrained/cct_14_7x2_224_imagenet.pth",
+    "cct_14_7x2_384": "http://ix.cs.uoregon.edu/~alih/compact-transformers/checkpoints/finetuned/cct_14_7x2_384_imagenet.pth",
+    "cct_14_7x2_384_fl": "http://ix.cs.uoregon.edu/~alih/compact-transformers/checkpoints/finetuned/cct_14_7x2_384_flowers102.pth",
+}
+
+
+class CCT(nn.Module):
+    def __init__(
+        self,
+        img_size=224,
+        embedding_dim=768,
+        n_input_channels=3,
+        n_conv_layers=1,
+        kernel_size=7,
+        stride=2,
+        padding=3,
+        pooling_kernel_size=3,
+        pooling_stride=2,
+        pooling_padding=1,
+        dropout=0.0,
+        attention_dropout=0.1,
+        stochastic_depth=0.1,
+        num_layers=14,
+        num_heads=6,
+        mlp_ratio=4.0,
+        num_classes=1000,
+        positional_embedding="learnable",
+        *args,
+        **kwargs,
+    ):
+        super(CCT, self).__init__()
+
+        self.tokenizer = Tokenizer(
+            n_input_channels=n_input_channels,
+            n_output_channels=embedding_dim,
+            kernel_size=kernel_size,
+            stride=stride,
+            padding=padding,
+            pooling_kernel_size=pooling_kernel_size,
+            pooling_stride=pooling_stride,
+            pooling_padding=pooling_padding,
+            max_pool=True,
+            activation=nn.ReLU,
+            n_conv_layers=n_conv_layers,
+            conv_bias=False,
+        )
+
+        self.classifier = TransformerClassifier(
+            sequence_length=self.tokenizer.sequence_length(
+                n_channels=n_input_channels, height=img_size, width=img_size
+            ),
+            embedding_dim=embedding_dim,
+            seq_pool=True,
+            dropout=dropout,
+            attention_dropout=attention_dropout,
+            stochastic_depth=stochastic_depth,
+            num_layers=num_layers,
+            num_heads=num_heads,
+            mlp_ratio=mlp_ratio,
+            num_classes=num_classes,
+            positional_embedding=positional_embedding,
+        )
+
+    def forward(self, x):
+        x = self.tokenizer(x)
+        return self.classifier(x)
+
+
+def _cct(
+    arch,
+    pretrained,
+    progress,
+    num_layers,
+    num_heads,
+    mlp_ratio,
+    embedding_dim,
+    kernel_size=3,
+    stride=None,
+    padding=None,
+    positional_embedding="learnable",
+    *args,
+    **kwargs,
+):
+    stride = stride if stride is not None else max(1, (kernel_size // 2) - 1)
+    padding = padding if padding is not None else max(1, (kernel_size // 2))
+    model = CCT(
+        num_layers=num_layers,
+        num_heads=num_heads,
+        mlp_ratio=mlp_ratio,
+        embedding_dim=embedding_dim,
+        kernel_size=kernel_size,
+        stride=stride,
+        padding=padding,
+        *args,
+        **kwargs,
+    )
+
+    if pretrained:
+        if arch in model_urls:
+            state_dict = load_state_dict_from_url(model_urls[arch], progress=progress)
+            if positional_embedding == "learnable":
+                state_dict = pe_check(model, state_dict)
+            elif positional_embedding == "sine":
+                state_dict["classifier.positional_emb"] = model.state_dict()[
+                    "classifier.positional_emb"
+                ]
+            state_dict = fc_check(model, state_dict)
+            model.load_state_dict(state_dict)
+        else:
+            raise RuntimeError(f"Variant {arch} does not yet have pretrained weights.")
+    return model
+
+
+def cct_2(arch, pretrained, progress, *args, **kwargs):
+    return _cct(
+        arch,
+        pretrained,
+        progress,
+        num_layers=2,
+        num_heads=2,
+        mlp_ratio=1,
+        embedding_dim=128,
+        *args,
+        **kwargs,
+    )
+
+
+def cct_4(arch, pretrained, progress, *args, **kwargs):
+    return _cct(
+        arch,
+        pretrained,
+        progress,
+        num_layers=4,
+        num_heads=2,
+        mlp_ratio=1,
+        embedding_dim=128,
+        *args,
+        **kwargs,
+    )
+
+
+def cct_6(arch, pretrained, progress, *args, **kwargs):
+    return _cct(
+        arch,
+        pretrained,
+        progress,
+        num_layers=6,
+        num_heads=4,
+        mlp_ratio=2,
+        embedding_dim=256,
+        *args,
+        **kwargs,
+    )
+
+
+def cct_7(arch, pretrained, progress, *args, **kwargs):
+    return _cct(
+        arch,
+        pretrained,
+        progress,
+        num_layers=7,
+        num_heads=4,
+        mlp_ratio=2,
+        embedding_dim=256,
+        *args,
+        **kwargs,
+    )
+
+
+def cct_14(arch, pretrained, progress, *args, **kwargs):
+    return _cct(
+        arch,
+        pretrained,
+        progress,
+        num_layers=14,
+        num_heads=6,
+        mlp_ratio=3,
+        embedding_dim=384,
+        *args,
+        **kwargs,
+    )
+
+
+@register_model
+def cct_2_3x2_32(
+    pretrained=False,
+    progress=False,
+    img_size=32,
+    positional_embedding="learnable",
+    num_classes=10,
+    *args,
+    **kwargs,
+):
+    return cct_2(
+        "cct_2_3x2_32",
+        pretrained,
+        progress,
+        kernel_size=3,
+        n_conv_layers=2,
+        img_size=img_size,
+        positional_embedding=positional_embedding,
+        num_classes=num_classes,
+        *args,
+        **kwargs,
+    )
+
+
+@register_model
+def cct_2_3x2_32_sine(
+    pretrained=False,
+    progress=False,
+    img_size=32,
+    positional_embedding="sine",
+    num_classes=10,
+    *args,
+    **kwargs,
+):
+    return cct_2(
+        "cct_2_3x2_32_sine",
+        pretrained,
+        progress,
+        kernel_size=3,
+        n_conv_layers=2,
+        img_size=img_size,
+        positional_embedding=positional_embedding,
+        num_classes=num_classes,
+        *args,
+        **kwargs,
+    )
+
+
+@register_model
+def cct_4_3x2_32(
+    pretrained=False,
+    progress=False,
+    img_size=32,
+    positional_embedding="learnable",
+    num_classes=10,
+    *args,
+    **kwargs,
+):
+    return cct_4(
+        "cct_4_3x2_32",
+        pretrained,
+        progress,
+        kernel_size=3,
+        n_conv_layers=2,
+        img_size=img_size,
+        positional_embedding=positional_embedding,
+        num_classes=num_classes,
+        *args,
+        **kwargs,
+    )
+
+
+@register_model
+def cct_4_3x2_32_sine(
+    pretrained=False,
+    progress=False,
+    img_size=32,
+    positional_embedding="sine",
+    num_classes=10,
+    *args,
+    **kwargs,
+):
+    return cct_4(
+        "cct_4_3x2_32_sine",
+        pretrained,
+        progress,
+        kernel_size=3,
+        n_conv_layers=2,
+        img_size=img_size,
+        positional_embedding=positional_embedding,
+        num_classes=num_classes,
+        *args,
+        **kwargs,
+    )
+
+
+@register_model
+def cct_6_3x1_32(
+    pretrained=False,
+    progress=False,
+    img_size=32,
+    positional_embedding="learnable",
+    num_classes=10,
+    *args,
+    **kwargs,
+):
+    return cct_6(
+        "cct_6_3x1_32",
+        pretrained,
+        progress,
+        kernel_size=3,
+        n_conv_layers=1,
+        img_size=img_size,
+        positional_embedding=positional_embedding,
+        num_classes=num_classes,
+        *args,
+        **kwargs,
+    )
+
+
+@register_model
+def cct_6_3x1_32_sine(
+    pretrained=False,
+    progress=False,
+    img_size=32,
+    positional_embedding="sine",
+    num_classes=10,
+    *args,
+    **kwargs,
+):
+    return cct_6(
+        "cct_6_3x1_32_sine",
+        pretrained,
+        progress,
+        kernel_size=3,
+        n_conv_layers=1,
+        img_size=img_size,
+        positional_embedding=positional_embedding,
+        num_classes=num_classes,
+        *args,
+        **kwargs,
+    )
+
+
+@register_model
+def cct_6_3x2_32(
+    pretrained=False,
+    progress=False,
+    img_size=32,
+    positional_embedding="learnable",
+    num_classes=10,
+    *args,
+    **kwargs,
+):
+    return cct_6(
+        "cct_6_3x2_32",
+        pretrained,
+        progress,
+        kernel_size=3,
+        n_conv_layers=2,
+        img_size=img_size,
+        positional_embedding=positional_embedding,
+        num_classes=num_classes,
+        *args,
+        **kwargs,
+    )
+
+
+@register_model
+def cct_6_3x2_32_sine(
+    pretrained=False,
+    progress=False,
+    img_size=32,
+    positional_embedding="sine",
+    num_classes=10,
+    *args,
+    **kwargs,
+):
+    return cct_6(
+        "cct_6_3x2_32_sine",
+        pretrained,
+        progress,
+        kernel_size=3,
+        n_conv_layers=2,
+        img_size=img_size,
+        positional_embedding=positional_embedding,
+        num_classes=num_classes,
+        *args,
+        **kwargs,
+    )
+
+
+@register_model
+def cct_7_3x1_32(
+    pretrained=False,
+    progress=False,
+    img_size=32,
+    positional_embedding="learnable",
+    num_classes=10,
+    *args,
+    **kwargs,
+):
+    return cct_7(
+        "cct_7_3x1_32",
+        pretrained,
+        progress,
+        kernel_size=3,
+        n_conv_layers=1,
+        img_size=img_size,
+        positional_embedding=positional_embedding,
+        num_classes=num_classes,
+        *args,
+        **kwargs,
+    )
+
+
+@register_model
+def cct_7_3x1_32_sine(
+    pretrained=False,
+    progress=False,
+    img_size=32,
+    positional_embedding="sine",
+    num_classes=10,
+    *args,
+    **kwargs,
+):
+    return cct_7(
+        "cct_7_3x1_32_sine",
+        pretrained,
+        progress,
+        kernel_size=3,
+        n_conv_layers=1,
+        img_size=img_size,
+        positional_embedding=positional_embedding,
+        num_classes=num_classes,
+        *args,
+        **kwargs,
+    )
+
+
+@register_model
+def cct_7_3x1_32_c100(
+    pretrained=False,
+    progress=False,
+    img_size=32,
+    positional_embedding="learnable",
+    num_classes=100,
+    *args,
+    **kwargs,
+):
+    return cct_7(
+        "cct_7_3x1_32_c100",
+        pretrained,
+        progress,
+        kernel_size=3,
+        n_conv_layers=1,
+        img_size=img_size,
+        positional_embedding=positional_embedding,
+        num_classes=num_classes,
+        *args,
+        **kwargs,
+    )
+
+
+@register_model
+def cct_7_3x1_32_sine_c100(
+    pretrained=False,
+    progress=False,
+    img_size=32,
+    positional_embedding="sine",
+    num_classes=100,
+    *args,
+    **kwargs,
+):
+    return cct_7(
+        "cct_7_3x1_32_sine_c100",
+        pretrained,
+        progress,
+        kernel_size=3,
+        n_conv_layers=1,
+        img_size=img_size,
+        positional_embedding=positional_embedding,
+        num_classes=num_classes,
+        *args,
+        **kwargs,
+    )
+
+
+@register_model
+def cct_7_3x2_32(
+    pretrained=False,
+    progress=False,
+    img_size=32,
+    positional_embedding="learnable",
+    num_classes=10,
+    *args,
+    **kwargs,
+):
+    return cct_7(
+        "cct_7_3x2_32",
+        pretrained,
+        progress,
+        kernel_size=3,
+        n_conv_layers=2,
+        img_size=img_size,
+        positional_embedding=positional_embedding,
+        num_classes=num_classes,
+        *args,
+        **kwargs,
+    )
+
+
+@register_model
+def cct_7_3x2_32_sine(
+    pretrained=False,
+    progress=False,
+    img_size=32,
+    positional_embedding="sine",
+    num_classes=10,
+    *args,
+    **kwargs,
+):
+    return cct_7(
+        "cct_7_3x2_32_sine",
+        pretrained,
+        progress,
+        kernel_size=3,
+        n_conv_layers=2,
+        img_size=img_size,
+        positional_embedding=positional_embedding,
+        num_classes=num_classes,
+        *args,
+        **kwargs,
+    )
+
+
+@register_model
+def cct_7_7x2_224(
+    pretrained=False,
+    progress=False,
+    img_size=224,
+    positional_embedding="learnable",
+    num_classes=102,
+    *args,
+    **kwargs,
+):
+    return cct_7(
+        "cct_7_7x2_224",
+        pretrained,
+        progress,
+        kernel_size=7,
+        n_conv_layers=2,
+        img_size=img_size,
+        positional_embedding=positional_embedding,
+        num_classes=num_classes,
+        *args,
+        **kwargs,
+    )
+
+
+@register_model
+def cct_7_7x2_224_sine(
+    pretrained=False,
+    progress=False,
+    img_size=224,
+    positional_embedding="sine",
+    num_classes=102,
+    *args,
+    **kwargs,
+):
+    return cct_7(
+        "cct_7_7x2_224_sine",
+        pretrained,
+        progress,
+        kernel_size=7,
+        n_conv_layers=2,
+        img_size=img_size,
+        positional_embedding=positional_embedding,
+        num_classes=num_classes,
+        *args,
+        **kwargs,
+    )
+
+
+@register_model
+def cct_14_7x2_224(
+    pretrained=False,
+    progress=False,
+    img_size=224,
+    positional_embedding="learnable",
+    num_classes=1000,
+    *args,
+    **kwargs,
+):
+    return cct_14(
+        "cct_14_7x2_224",
+        pretrained,
+        progress,
+        kernel_size=7,
+        n_conv_layers=2,
+        img_size=img_size,
+        positional_embedding=positional_embedding,
+        num_classes=num_classes,
+        *args,
+        **kwargs,
+    )
+
+
+@register_model
+def cct_14_7x2_384(
+    pretrained=False,
+    progress=False,
+    img_size=384,
+    positional_embedding="learnable",
+    num_classes=1000,
+    *args,
+    **kwargs,
+):
+    return cct_14(
+        "cct_14_7x2_384",
+        pretrained,
+        progress,
+        kernel_size=7,
+        n_conv_layers=2,
+        img_size=img_size,
+        positional_embedding=positional_embedding,
+        num_classes=num_classes,
+        *args,
+        **kwargs,
+    )
+
+
+@register_model
+def cct_14_7x2_384_fl(
+    pretrained=False,
+    progress=False,
+    img_size=384,
+    positional_embedding="learnable",
+    num_classes=102,
+    *args,
+    **kwargs,
+):
+    return cct_14(
+        "cct_14_7x2_384_fl",
+        pretrained,
+        progress,
+        kernel_size=7,
+        n_conv_layers=2,
+        img_size=img_size,
+        positional_embedding=positional_embedding,
+        num_classes=num_classes,
+        *args,
+        **kwargs,
+    )
diff --git a/src/blades/models/backbones/cctnets/cvt.py b/src/blades/models/backbones/cctnets/cvt.py
new file mode 100644
index 0000000..d4a74e6
--- /dev/null
+++ b/src/blades/models/backbones/cctnets/cvt.py
@@ -0,0 +1,318 @@
+import torch.nn as nn
+from torch.hub import load_state_dict_from_url
+
+from .utils.helpers import pe_check
+from .utils.tokenizer import Tokenizer
+from .utils.transformers import TransformerClassifier
+
+try:
+    from timm.models.registry import register_model
+except ImportError:
+    from .registry import register_model
+
+model_urls = {}
+
+
+class CVT(nn.Module):
+    def __init__(
+        self,
+        img_size=224,
+        embedding_dim=768,
+        n_input_channels=3,
+        kernel_size=16,
+        dropout=0.0,
+        attention_dropout=0.1,
+        stochastic_depth=0.1,
+        num_layers=14,
+        num_heads=6,
+        mlp_ratio=4.0,
+        num_classes=1000,
+        positional_embedding="learnable",
+        *args,
+        **kwargs,
+    ):
+        super(CVT, self).__init__()
+        assert img_size % kernel_size == 0, (
+            f"Image size ({img_size}) has to be"
+            f"divisible by patch size ({kernel_size})"
+        )
+        self.tokenizer = Tokenizer(
+            n_input_channels=n_input_channels,
+            n_output_channels=embedding_dim,
+            kernel_size=kernel_size,
+            stride=kernel_size,
+            padding=0,
+            max_pool=False,
+            activation=None,
+            n_conv_layers=1,
+            conv_bias=True,
+        )
+
+        self.classifier = TransformerClassifier(
+            sequence_length=self.tokenizer.sequence_length(
+                n_channels=n_input_channels, height=img_size, width=img_size
+            ),
+            embedding_dim=embedding_dim,
+            seq_pool=True,
+            dropout=dropout,
+            attention_dropout=attention_dropout,
+            stochastic_depth=stochastic_depth,
+            num_layers=num_layers,
+            num_heads=num_heads,
+            mlp_ratio=mlp_ratio,
+            num_classes=num_classes,
+            positional_embedding=positional_embedding,
+        )
+
+    def forward(self, x):
+        x = self.tokenizer(x)
+        return self.classifier(x)
+
+
+def _cvt(
+    arch,
+    pretrained,
+    progress,
+    num_layers,
+    num_heads,
+    mlp_ratio,
+    embedding_dim,
+    kernel_size=4,
+    positional_embedding="learnable",
+    *args,
+    **kwargs,
+):
+    model = CVT(
+        num_layers=num_layers,
+        num_heads=num_heads,
+        mlp_ratio=mlp_ratio,
+        embedding_dim=embedding_dim,
+        kernel_size=kernel_size,
+        *args,
+        **kwargs,
+    )
+
+    if pretrained and arch in model_urls:
+        state_dict = load_state_dict_from_url(model_urls[arch], progress=progress)
+        if positional_embedding == "learnable":
+            state_dict = pe_check(model, state_dict)
+        elif positional_embedding == "sine":
+            state_dict["classifier.positional_emb"] = model.state_dict()[
+                "classifier.positional_emb"
+            ]
+        model.load_state_dict(state_dict)
+    return model
+
+
+def cvt_2(*args, **kwargs):
+    return _cvt(
+        num_layers=2, num_heads=2, mlp_ratio=1, embedding_dim=128, *args, **kwargs
+    )
+
+
+def cvt_4(*args, **kwargs):
+    return _cvt(
+        num_layers=4, num_heads=2, mlp_ratio=1, embedding_dim=128, *args, **kwargs
+    )
+
+
+def cvt_6(*args, **kwargs):
+    return _cvt(
+        num_layers=6, num_heads=4, mlp_ratio=2, embedding_dim=256, *args, **kwargs
+    )
+
+
+def cvt_7(*args, **kwargs):
+    return _cvt(
+        num_layers=7, num_heads=4, mlp_ratio=2, embedding_dim=256, *args, **kwargs
+    )
+
+
+def cvt_8(*args, **kwargs):
+    return _cvt(
+        num_layers=8, num_heads=4, mlp_ratio=2, embedding_dim=256, *args, **kwargs
+    )
+
+
+@register_model
+def cvt_2_4_32(
+    pretrained=False,
+    progress=False,
+    img_size=32,
+    positional_embedding="learnable",
+    num_classes=10,
+    *args,
+    **kwargs,
+):
+    return cvt_2(
+        "cvt_2_4_32",
+        pretrained,
+        progress,
+        kernel_size=4,
+        img_size=img_size,
+        positional_embedding=positional_embedding,
+        num_classes=num_classes,
+        *args,
+        **kwargs,
+    )
+
+
+@register_model
+def cvt_2_4_32_sine(
+    pretrained=False,
+    progress=False,
+    img_size=32,
+    positional_embedding="sine",
+    num_classes=10,
+    *args,
+    **kwargs,
+):
+    return cvt_2(
+        "cvt_2_4_32_sine",
+        pretrained,
+        progress,
+        kernel_size=4,
+        img_size=img_size,
+        positional_embedding=positional_embedding,
+        num_classes=num_classes,
+        *args,
+        **kwargs,
+    )
+
+
+@register_model
+def cvt_4_4_32(
+    pretrained=False,
+    progress=False,
+    img_size=32,
+    positional_embedding="learnable",
+    num_classes=10,
+    *args,
+    **kwargs,
+):
+    return cvt_4(
+        "cvt_4_4_32",
+        pretrained,
+        progress,
+        kernel_size=4,
+        img_size=img_size,
+        positional_embedding=positional_embedding,
+        num_classes=num_classes,
+        *args,
+        **kwargs,
+    )
+
+
+@register_model
+def cvt_4_4_32_sine(
+    pretrained=False,
+    progress=False,
+    img_size=32,
+    positional_embedding="sine",
+    num_classes=10,
+    *args,
+    **kwargs,
+):
+    return cvt_4(
+        "cvt_4_4_32_sine",
+        pretrained,
+        progress,
+        kernel_size=4,
+        img_size=img_size,
+        positional_embedding=positional_embedding,
+        num_classes=num_classes,
+        *args,
+        **kwargs,
+    )
+
+
+@register_model
+def cvt_6_4_32(
+    pretrained=False,
+    progress=False,
+    img_size=32,
+    positional_embedding="learnable",
+    num_classes=10,
+    *args,
+    **kwargs,
+):
+    return cvt_6(
+        "cvt_6_4_32",
+        pretrained,
+        progress,
+        kernel_size=4,
+        img_size=img_size,
+        positional_embedding=positional_embedding,
+        num_classes=num_classes,
+        *args,
+        **kwargs,
+    )
+
+
+@register_model
+def cvt_6_4_32_sine(
+    pretrained=False,
+    progress=False,
+    img_size=32,
+    positional_embedding="sine",
+    num_classes=10,
+    *args,
+    **kwargs,
+):
+    return cvt_6(
+        "cvt_6_4_32_sine",
+        pretrained,
+        progress,
+        kernel_size=4,
+        img_size=img_size,
+        positional_embedding=positional_embedding,
+        num_classes=num_classes,
+        *args,
+        **kwargs,
+    )
+
+
+@register_model
+def cvt_7_4_32(
+    pretrained=False,
+    progress=False,
+    img_size=32,
+    positional_embedding="learnable",
+    num_classes=10,
+    *args,
+    **kwargs,
+):
+    return cvt_7(
+        "cvt_7_4_32",
+        pretrained,
+        progress,
+        kernel_size=4,
+        img_size=img_size,
+        positional_embedding=positional_embedding,
+        num_classes=num_classes,
+        *args,
+        **kwargs,
+    )
+
+
+@register_model
+def cvt_7_4_32_sine(
+    pretrained=False,
+    progress=False,
+    img_size=32,
+    positional_embedding="sine",
+    num_classes=10,
+    *args,
+    **kwargs,
+):
+    return cvt_7(
+        "cvt_7_4_32_sine",
+        pretrained,
+        progress,
+        kernel_size=4,
+        img_size=img_size,
+        positional_embedding=positional_embedding,
+        num_classes=num_classes,
+        *args,
+        **kwargs,
+    )
diff --git a/src/blades/models/cifar10/cctnets/registry.py b/src/blades/models/backbones/cctnets/registry.py
similarity index 100%
rename from src/blades/models/cifar10/cctnets/registry.py
rename to src/blades/models/backbones/cctnets/registry.py
diff --git a/src/blades/models/backbones/cctnets/utils/__init__.py b/src/blades/models/backbones/cctnets/utils/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/src/blades/models/backbones/cctnets/utils/embedder.py b/src/blades/models/backbones/cctnets/utils/embedder.py
new file mode 100644
index 0000000..b1a9c83
--- /dev/null
+++ b/src/blades/models/backbones/cctnets/utils/embedder.py
@@ -0,0 +1,50 @@
+import torch.nn as nn
+
+
+class Embedder(nn.Module):
+    def __init__(
+        self,
+        word_embedding_dim=300,
+        vocab_size=100000,
+        padding_idx=1,
+        pretrained_weight=None,
+        embed_freeze=False,
+        *args,
+        **kwargs
+    ):
+        super(Embedder, self).__init__()
+        self.embeddings = (
+            nn.Embedding.from_pretrained(pretrained_weight, freeze=embed_freeze)
+            if pretrained_weight is not None
+            else nn.Embedding(
+                vocab_size,
+                word_embedding_dim,
+                padding_idx=padding_idx,
+            )
+        )
+        self.embeddings.weight.requires_grad = not embed_freeze
+
+    def forward_mask(self, mask):
+        bsz, seq_len = mask.shape
+        new_mask = mask.view(bsz, seq_len, 1)
+        new_mask = new_mask.sum(-1)
+        new_mask = new_mask > 0
+        return new_mask
+
+    def forward(self, x, mask=None):
+        embed = self.embeddings(x)
+        embed = (
+            embed
+            if mask is None
+            else embed * self.forward_mask(mask).unsqueeze(-1).float()
+        )
+        return embed, mask
+
+    @staticmethod
+    def init_weight(m):
+        if isinstance(m, nn.Linear):
+            nn.init.trunc_normal_(m.weight, std=0.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        else:
+            nn.init.normal_(m.weight)
diff --git a/src/blades/models/cifar10/cctnets/utils/helpers.py b/src/blades/models/backbones/cctnets/utils/helpers.py
similarity index 62%
rename from src/blades/models/cifar10/cctnets/utils/helpers.py
rename to src/blades/models/backbones/cctnets/utils/helpers.py
index c3e34c4..9a314bc 100644
--- a/src/blades/models/cifar10/cctnets/utils/helpers.py
+++ b/src/blades/models/backbones/cctnets/utils/helpers.py
@@ -4,7 +4,7 @@
 import torch
 import torch.nn.functional as F
 
-_logger = logging.getLogger('train')
+_logger = logging.getLogger("train")
 
 
 def resize_pos_embed(posemb, posemb_new, num_tokens=1):
@@ -21,25 +21,35 @@ def resize_pos_embed(posemb, posemb_new, num_tokens=1):
     gs_old = int(math.sqrt(len(posemb_grid)))
     gs_new = int(math.sqrt(ntok_new))
     posemb_grid = posemb_grid.reshape(1, gs_old, gs_old, -1).permute(0, 3, 1, 2)
-    posemb_grid = F.interpolate(posemb_grid, size=(gs_new, gs_new), mode='bilinear')
+    posemb_grid = F.interpolate(posemb_grid, size=(gs_new, gs_new), mode="bilinear")
     posemb_grid = posemb_grid.permute(0, 2, 3, 1).reshape(1, gs_new * gs_new, -1)
     posemb = torch.cat([posemb_tok, posemb_grid], dim=1)
     return posemb
 
 
-def pe_check(model, state_dict, pe_key='classifier.positional_emb'):
-    if pe_key is not None and pe_key in state_dict.keys() and pe_key in model.state_dict().keys():
+def pe_check(model, state_dict, pe_key="classifier.positional_emb"):
+    if (
+        pe_key is not None
+        and pe_key in state_dict.keys()
+        and pe_key in model.state_dict().keys()
+    ):
         if model.state_dict()[pe_key].shape != state_dict[pe_key].shape:
-            state_dict[pe_key] = resize_pos_embed(state_dict[pe_key],
-                                                  model.state_dict()[pe_key],
-                                                  num_tokens=model.classifier.num_tokens)
+            state_dict[pe_key] = resize_pos_embed(
+                state_dict[pe_key],
+                model.state_dict()[pe_key],
+                num_tokens=model.classifier.num_tokens,
+            )
     return state_dict
 
 
-def fc_check(model, state_dict, fc_key='classifier.fc'):
-    for key in [f'{fc_key}.weight', f'{fc_key}.bias']:
-        if key is not None and key in state_dict.keys() and key in model.state_dict().keys():
+def fc_check(model, state_dict, fc_key="classifier.fc"):
+    for key in [f"{fc_key}.weight", f"{fc_key}.bias"]:
+        if (
+            key is not None
+            and key in state_dict.keys()
+            and key in model.state_dict().keys()
+        ):
             if model.state_dict()[key].shape != state_dict[key].shape:
-                _logger.warning(f'Removing {key}, number of classes has changed.')
+                _logger.warning(f"Removing {key}, number of classes has changed.")
                 state_dict[key] = model.state_dict()[key]
     return state_dict
diff --git a/src/blades/models/cifar10/cctnets/utils/stochastic_depth.py b/src/blades/models/backbones/cctnets/utils/stochastic_depth.py
similarity index 65%
rename from src/blades/models/cifar10/cctnets/utils/stochastic_depth.py
rename to src/blades/models/backbones/cctnets/utils/stochastic_depth.py
index 51c12d4..d3d7b59 100644
--- a/src/blades/models/cifar10/cctnets/utils/stochastic_depth.py
+++ b/src/blades/models/backbones/cctnets/utils/stochastic_depth.py
@@ -5,20 +5,23 @@
 import torch.nn as nn
 
 
-def drop_path(x, drop_prob: float = 0., training: bool = False):
-    """
-    Obtained from: github.com:rwightman/pytorch-image-settings
-    Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
+def drop_path(x, drop_prob: float = 0.0, training: bool = False):
+    """Obtained from: github.com:rwightman/pytorch-image-settings Drop paths
+    (Stochastic Depth) per sample (when applied in main path of residual
+    blocks).
+
     This is the same as the DropConnect impl I created for EfficientNet, etc networks, however,
     the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...
     See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for
     changing the layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use
     'survival rate' as the argument.
     """
-    if drop_prob == 0. or not training:
+    if drop_prob == 0.0 or not training:
         return x
     keep_prob = 1 - drop_prob
-    shape = (x.shape[0],) + (1,) * (x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
+    shape = (x.shape[0],) + (1,) * (
+        x.ndim - 1
+    )  # work with diff dim tensors, not just 2D ConvNets
     random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
     random_tensor.floor_()  # binarize
     output = x.div(keep_prob) * random_tensor
@@ -26,14 +29,13 @@ def drop_path(x, drop_prob: float = 0., training: bool = False):
 
 
 class DropPath(nn.Module):
-    """
-    Obtained from: github.com:rwightman/pytorch-image-settings
-    Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks).
-    """
-    
+    """Obtained from: github.com:rwightman/pytorch-image-settings Drop paths
+    (Stochastic Depth) per sample  (when applied in main path of residual
+    blocks)."""
+
     def __init__(self, drop_prob=None):
         super(DropPath, self).__init__()
         self.drop_prob = drop_prob
-    
+
     def forward(self, x):
         return drop_path(x, self.drop_prob, self.training)
diff --git a/src/blades/models/backbones/cctnets/utils/tokenizer.py b/src/blades/models/backbones/cctnets/utils/tokenizer.py
new file mode 100644
index 0000000..3e06dd2
--- /dev/null
+++ b/src/blades/models/backbones/cctnets/utils/tokenizer.py
@@ -0,0 +1,155 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class Tokenizer(nn.Module):
+    def __init__(
+        self,
+        kernel_size,
+        stride,
+        padding,
+        pooling_kernel_size=3,
+        pooling_stride=2,
+        pooling_padding=1,
+        n_conv_layers=1,
+        n_input_channels=3,
+        n_output_channels=64,
+        in_planes=64,
+        activation=None,
+        max_pool=True,
+        conv_bias=False,
+    ):
+        super(Tokenizer, self).__init__()
+
+        n_filter_list = (
+            [n_input_channels]
+            + [in_planes for _ in range(n_conv_layers - 1)]
+            + [n_output_channels]
+        )
+
+        self.conv_layers = nn.Sequential(
+            *[
+                nn.Sequential(
+                    nn.Conv2d(
+                        n_filter_list[i],
+                        n_filter_list[i + 1],
+                        kernel_size=(kernel_size, kernel_size),
+                        stride=(stride, stride),
+                        padding=(padding, padding),
+                        bias=conv_bias,
+                    ),
+                    nn.Identity() if activation is None else activation(),
+                    nn.MaxPool2d(
+                        kernel_size=pooling_kernel_size,
+                        stride=pooling_stride,
+                        padding=pooling_padding,
+                    )
+                    if max_pool
+                    else nn.Identity(),
+                )
+                for i in range(n_conv_layers)
+            ]
+        )
+
+        self.flattener = nn.Flatten(2, 3)
+        self.apply(self.init_weight)
+
+    def sequence_length(self, n_channels=3, height=224, width=224):
+        return self.forward(torch.zeros((1, n_channels, height, width))).shape[1]
+
+    def forward(self, x):
+        return self.flattener(self.conv_layers(x)).transpose(-2, -1)
+
+    @staticmethod
+    def init_weight(m):
+        if isinstance(m, nn.Conv2d):
+            nn.init.kaiming_normal_(m.weight)
+
+
+class TextTokenizer(nn.Module):
+    def __init__(
+        self,
+        kernel_size,
+        stride,
+        padding,
+        pooling_kernel_size=3,
+        pooling_stride=2,
+        pooling_padding=1,
+        embedding_dim=300,
+        n_output_channels=128,
+        activation=None,
+        max_pool=True,
+        *args,
+        **kwargs
+    ):
+        super(TextTokenizer, self).__init__()
+
+        self.max_pool = max_pool
+        self.conv_layers = nn.Sequential(
+            nn.Conv2d(
+                1,
+                n_output_channels,
+                kernel_size=(kernel_size, embedding_dim),
+                stride=(stride, 1),
+                padding=(padding, 0),
+                bias=False,
+            ),
+            nn.Identity() if activation is None else activation(),
+            nn.MaxPool2d(
+                kernel_size=(pooling_kernel_size, 1),
+                stride=(pooling_stride, 1),
+                padding=(pooling_padding, 0),
+            )
+            if max_pool
+            else nn.Identity(),
+        )
+
+        self.apply(self.init_weight)
+
+    def seq_len(self, seq_len=32, embed_dim=300):
+        return self.forward(torch.zeros((1, seq_len, embed_dim)))[0].shape[1]
+
+    def forward_mask(self, mask):
+        new_mask = mask.unsqueeze(1).float()
+        cnn_weight = torch.ones(
+            (1, 1, self.conv_layers[0].kernel_size[0]),
+            device=mask.device,
+            dtype=torch.float,
+        )
+        new_mask = F.conv1d(
+            new_mask,
+            cnn_weight,
+            None,
+            self.conv_layers[0].stride[0],
+            self.conv_layers[0].padding[0],
+            1,
+            1,
+        )
+        if self.max_pool:
+            new_mask = F.max_pool1d(
+                new_mask,
+                self.conv_layers[2].kernel_size[0],
+                self.conv_layers[2].stride[0],
+                self.conv_layers[2].padding[0],
+                1,
+                False,
+                False,
+            )
+        new_mask = new_mask.squeeze(1)
+        new_mask = new_mask > 0
+        return new_mask
+
+    def forward(self, x, mask=None):
+        x = x.unsqueeze(1)
+        x = self.conv_layers(x)
+        x = x.transpose(1, 3).squeeze(1)
+        if mask is not None:
+            mask = self.forward_mask(mask).unsqueeze(-1).float()
+            x = x * mask
+        return x, mask
+
+    @staticmethod
+    def init_weight(m):
+        if isinstance(m, nn.Conv2d):
+            nn.init.kaiming_normal_(m.weight)
diff --git a/src/blades/models/cifar10/cctnets/utils/transformers.py b/src/blades/models/backbones/cctnets/utils/transformers.py
similarity index 56%
rename from src/blades/models/cifar10/cctnets/utils/transformers.py
rename to src/blades/models/backbones/cctnets/utils/transformers.py
index 1310e0a..0a1b137 100644
--- a/src/blades/models/cifar10/cctnets/utils/transformers.py
+++ b/src/blades/models/backbones/cctnets/utils/transformers.py
@@ -1,35 +1,46 @@
 import torch
 import torch.nn.functional as F
-from torch.nn import Module, ModuleList, Linear, Dropout, LayerNorm, Identity, Parameter, init
+from torch.nn import (
+    Dropout,
+    Identity,
+    LayerNorm,
+    Linear,
+    Module,
+    ModuleList,
+    Parameter,
+    init,
+)
 
 from .stochastic_depth import DropPath
 
 
 class Attention(Module):
-    """
-    Obtained from timm: github.com:rwightman/pytorch-image-settings
-    """
-    
+    """Obtained from timm: github.com:rwightman/pytorch-image-settings."""
+
     def __init__(self, dim, num_heads=8, attention_dropout=0.1, projection_dropout=0.1):
         super().__init__()
         self.num_heads = num_heads
         head_dim = dim // self.num_heads
-        self.scale = head_dim ** -0.5
-        
+        self.scale = head_dim**-0.5
+
         self.qkv = Linear(dim, dim * 3, bias=False)
         self.attn_drop = Dropout(attention_dropout)
         self.proj = Linear(dim, dim)
         self.proj_drop = Dropout(projection_dropout)
-    
+
     def forward(self, x):
         B, N, C = x.shape
-        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+        qkv = (
+            self.qkv(x)
+            .reshape(B, N, 3, self.num_heads, C // self.num_heads)
+            .permute(2, 0, 3, 1, 4)
+        )
         q, k, v = qkv[0], qkv[1], qkv[2]
-        
+
         attn = (q @ k.transpose(-2, -1)) * self.scale
         attn = attn.softmax(dim=-1)
         attn = self.attn_drop(attn)
-        
+
         x = (attn @ v).transpose(1, 2).reshape(B, N, C)
         x = self.proj(x)
         x = self.proj_drop(x)
@@ -41,30 +52,34 @@ def __init__(self, dim, num_heads=8, attention_dropout=0.1, projection_dropout=0
         super().__init__()
         self.num_heads = num_heads
         head_dim = dim // self.num_heads
-        self.scale = head_dim ** -0.5
-        
+        self.scale = head_dim**-0.5
+
         self.qkv = Linear(dim, dim * 3, bias=False)
         self.attn_drop = Dropout(attention_dropout)
         self.proj = Linear(dim, dim)
         self.proj_drop = Dropout(projection_dropout)
-    
+
     def forward(self, x, mask=None):
         B, N, C = x.shape
-        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+        qkv = (
+            self.qkv(x)
+            .reshape(B, N, 3, self.num_heads, C // self.num_heads)
+            .permute(2, 0, 3, 1, 4)
+        )
         q, k, v = qkv[0], qkv[1], qkv[2]
-        
+
         attn = (q @ k.transpose(-2, -1)) * self.scale
-        
+
         if mask is not None:
             mask_value = -torch.finfo(attn.dtype).max
-            assert mask.shape[-1] == attn.shape[-1], 'mask has incorrect dimensions'
+            assert mask.shape[-1] == attn.shape[-1], "mask has incorrect dimensions"
             mask = mask[:, None, :] * mask[:, :, None]
             mask = mask.unsqueeze(1).repeat(1, self.num_heads, 1, 1)
             attn.masked_fill_(~mask, mask_value)
-        
+
         attn = attn.softmax(dim=-1)
         attn = self.attn_drop(attn)
-        
+
         x = (attn @ v).transpose(1, 2).reshape(B, N, C)
         x = self.proj(x)
         x = self.proj_drop(x)
@@ -72,27 +87,36 @@ def forward(self, x, mask=None):
 
 
 class TransformerEncoderLayer(Module):
-    """
-    Inspired by torch.nn.TransformerEncoderLayer and timm.
-    """
-    
-    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1,
-                 attention_dropout=0.1, drop_path_rate=0.1):
+    """Inspired by torch.nn.TransformerEncoderLayer and timm."""
+
+    def __init__(
+        self,
+        d_model,
+        nhead,
+        dim_feedforward=2048,
+        dropout=0.1,
+        attention_dropout=0.1,
+        drop_path_rate=0.1,
+    ):
         super(TransformerEncoderLayer, self).__init__()
         self.pre_norm = LayerNorm(d_model)
-        self.self_attn = Attention(dim=d_model, num_heads=nhead,
-                                   attention_dropout=attention_dropout, projection_dropout=dropout)
-        
+        self.self_attn = Attention(
+            dim=d_model,
+            num_heads=nhead,
+            attention_dropout=attention_dropout,
+            projection_dropout=dropout,
+        )
+
         self.linear1 = Linear(d_model, dim_feedforward)
         self.dropout1 = Dropout(dropout)
         self.norm1 = LayerNorm(d_model)
         self.linear2 = Linear(dim_feedforward, d_model)
         self.dropout2 = Dropout(dropout)
-        
+
         self.drop_path = DropPath(drop_path_rate) if drop_path_rate > 0 else Identity()
-        
+
         self.activation = F.gelu
-    
+
     def forward(self, src: torch.Tensor, *args, **kwargs) -> torch.Tensor:
         src = src + self.drop_path(self.self_attn(self.pre_norm(src)))
         src = self.norm1(src)
@@ -102,27 +126,36 @@ def forward(self, src: torch.Tensor, *args, **kwargs) -> torch.Tensor:
 
 
 class MaskedTransformerEncoderLayer(Module):
-    """
-    Inspired by torch.nn.TransformerEncoderLayer and timm.
-    """
-    
-    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1,
-                 attention_dropout=0.1, drop_path_rate=0.1):
+    """Inspired by torch.nn.TransformerEncoderLayer and timm."""
+
+    def __init__(
+        self,
+        d_model,
+        nhead,
+        dim_feedforward=2048,
+        dropout=0.1,
+        attention_dropout=0.1,
+        drop_path_rate=0.1,
+    ):
         super(MaskedTransformerEncoderLayer, self).__init__()
         self.pre_norm = LayerNorm(d_model)
-        self.self_attn = MaskedAttention(dim=d_model, num_heads=nhead,
-                                         attention_dropout=attention_dropout, projection_dropout=dropout)
-        
+        self.self_attn = MaskedAttention(
+            dim=d_model,
+            num_heads=nhead,
+            attention_dropout=attention_dropout,
+            projection_dropout=dropout,
+        )
+
         self.linear1 = Linear(d_model, dim_feedforward)
         self.dropout1 = Dropout(dropout)
         self.norm1 = LayerNorm(d_model)
         self.linear2 = Linear(dim_feedforward, d_model)
         self.dropout2 = Dropout(dropout)
-        
+
         self.drop_path = DropPath(drop_path_rate) if drop_path_rate > 0 else Identity()
-        
+
         self.activation = F.gelu
-    
+
     def forward(self, src: torch.Tensor, mask=None, *args, **kwargs) -> torch.Tensor:
         src = src + self.drop_path(self.self_attn(self.pre_norm(src), mask))
         src = self.norm1(src)
@@ -132,209 +165,266 @@ def forward(self, src: torch.Tensor, mask=None, *args, **kwargs) -> torch.Tensor
 
 
 class TransformerClassifier(Module):
-    def __init__(self,
-                 seq_pool=True,
-                 embedding_dim=768,
-                 num_layers=12,
-                 num_heads=12,
-                 mlp_ratio=4.0,
-                 num_classes=1000,
-                 dropout=0.1,
-                 attention_dropout=0.1,
-                 stochastic_depth=0.1,
-                 positional_embedding='learnable',
-                 sequence_length=None):
+    def __init__(
+        self,
+        seq_pool=True,
+        embedding_dim=768,
+        num_layers=12,
+        num_heads=12,
+        mlp_ratio=4.0,
+        num_classes=1000,
+        dropout=0.1,
+        attention_dropout=0.1,
+        stochastic_depth=0.1,
+        positional_embedding="learnable",
+        sequence_length=None,
+    ):
         super().__init__()
-        positional_embedding = positional_embedding if \
-            positional_embedding in ['sine', 'learnable', 'none'] else 'sine'
+        positional_embedding = (
+            positional_embedding
+            if positional_embedding in ["sine", "learnable", "none"]
+            else "sine"
+        )
         dim_feedforward = int(embedding_dim * mlp_ratio)
         self.embedding_dim = embedding_dim
         self.sequence_length = sequence_length
         self.seq_pool = seq_pool
         self.num_tokens = 0
-        
-        assert sequence_length is not None or positional_embedding == 'none', \
-            f"Positional embedding is set to {positional_embedding} and" \
+
+        assert sequence_length is not None or positional_embedding == "none", (
+            f"Positional embedding is set to {positional_embedding} and"
             f" the sequence length was not specified."
-        
+        )
+
         if not seq_pool:
             sequence_length += 1
-            self.class_emb = Parameter(torch.zeros(1, 1, self.embedding_dim),
-                                       requires_grad=True)
+            self.class_emb = Parameter(
+                torch.zeros(1, 1, self.embedding_dim), requires_grad=True
+            )
             self.num_tokens = 1
         else:
             self.attention_pool = Linear(self.embedding_dim, 1)
-        
-        if positional_embedding != 'none':
-            if positional_embedding == 'learnable':
-                self.positional_emb = Parameter(torch.zeros(1, sequence_length, embedding_dim),
-                                                requires_grad=True)
+
+        if positional_embedding != "none":
+            if positional_embedding == "learnable":
+                self.positional_emb = Parameter(
+                    torch.zeros(1, sequence_length, embedding_dim), requires_grad=True
+                )
                 init.trunc_normal_(self.positional_emb, std=0.2)
             else:
-                self.positional_emb = Parameter(self.sinusoidal_embedding(sequence_length, embedding_dim),
-                                                requires_grad=False)
+                self.positional_emb = Parameter(
+                    self.sinusoidal_embedding(sequence_length, embedding_dim),
+                    requires_grad=False,
+                )
         else:
             self.positional_emb = None
-        
+
         self.dropout = Dropout(p=dropout)
         dpr = [x.item() for x in torch.linspace(0, stochastic_depth, num_layers)]
-        self.blocks = ModuleList([
-            TransformerEncoderLayer(d_model=embedding_dim, nhead=num_heads,
-                                    dim_feedforward=dim_feedforward, dropout=dropout,
-                                    attention_dropout=attention_dropout, drop_path_rate=dpr[i])
-            for i in range(num_layers)])
+        self.blocks = ModuleList(
+            [
+                TransformerEncoderLayer(
+                    d_model=embedding_dim,
+                    nhead=num_heads,
+                    dim_feedforward=dim_feedforward,
+                    dropout=dropout,
+                    attention_dropout=attention_dropout,
+                    drop_path_rate=dpr[i],
+                )
+                for i in range(num_layers)
+            ]
+        )
         self.norm = LayerNorm(embedding_dim)
-        
+
         self.fc = Linear(embedding_dim, num_classes)
         self.apply(self.init_weight)
-    
+
     def forward(self, x):
         if self.positional_emb is None and x.size(1) < self.sequence_length:
-            x = F.pad(x, (0, 0, 0, self.n_channels - x.size(1)), mode='constant', value=0)
-        
+            x = F.pad(
+                x, (0, 0, 0, self.n_channels - x.size(1)), mode="constant", value=0
+            )
+
         if not self.seq_pool:
             cls_token = self.class_emb.expand(x.shape[0], -1, -1)
             x = torch.cat((cls_token, x), dim=1)
-        
+
         if self.positional_emb is not None:
             x += self.positional_emb
-        
+
         x = self.dropout(x)
-        
+
         for blk in self.blocks:
             x = blk(x)
         x = self.norm(x)
-        
+
         if self.seq_pool:
-            x = torch.matmul(F.softmax(self.attention_pool(x), dim=1).transpose(-1, -2), x).squeeze(-2)
+            x = torch.matmul(
+                F.softmax(self.attention_pool(x), dim=1).transpose(-1, -2), x
+            ).squeeze(-2)
         else:
             x = x[:, 0]
-        
+
         x = self.fc(x)
         return x
-    
+
     @staticmethod
     def init_weight(m):
         if isinstance(m, Linear):
-            init.trunc_normal_(m.weight, std=.02)
+            init.trunc_normal_(m.weight, std=0.02)
             if isinstance(m, Linear) and m.bias is not None:
                 init.constant_(m.bias, 0)
         elif isinstance(m, LayerNorm):
             init.constant_(m.bias, 0)
             init.constant_(m.weight, 1.0)
-    
+
     @staticmethod
     def sinusoidal_embedding(n_channels, dim):
-        pe = torch.FloatTensor([[p / (10000 ** (2 * (i // 2) / dim)) for i in range(dim)]
-                                for p in range(n_channels)])
+        pe = torch.FloatTensor(
+            [
+                [p / (10000 ** (2 * (i // 2) / dim)) for i in range(dim)]
+                for p in range(n_channels)
+            ]
+        )
         pe[:, 0::2] = torch.sin(pe[:, 0::2])
         pe[:, 1::2] = torch.cos(pe[:, 1::2])
         return pe.unsqueeze(0)
 
 
 class MaskedTransformerClassifier(Module):
-    def __init__(self,
-                 seq_pool=True,
-                 embedding_dim=768,
-                 num_layers=12,
-                 num_heads=12,
-                 mlp_ratio=4.0,
-                 num_classes=1000,
-                 dropout=0.1,
-                 attention_dropout=0.1,
-                 stochastic_depth=0.1,
-                 positional_embedding='sine',
-                 seq_len=None,
-                 *args, **kwargs):
+    def __init__(
+        self,
+        seq_pool=True,
+        embedding_dim=768,
+        num_layers=12,
+        num_heads=12,
+        mlp_ratio=4.0,
+        num_classes=1000,
+        dropout=0.1,
+        attention_dropout=0.1,
+        stochastic_depth=0.1,
+        positional_embedding="sine",
+        seq_len=None,
+        *args,
+        **kwargs,
+    ):
         super().__init__()
-        positional_embedding = positional_embedding if \
-            positional_embedding in ['sine', 'learnable', 'none'] else 'sine'
+        positional_embedding = (
+            positional_embedding
+            if positional_embedding in ["sine", "learnable", "none"]
+            else "sine"
+        )
         dim_feedforward = int(embedding_dim * mlp_ratio)
         self.embedding_dim = embedding_dim
         self.seq_len = seq_len
         self.seq_pool = seq_pool
         self.num_tokens = 0
-        
-        assert seq_len is not None or positional_embedding == 'none', \
-            f"Positional embedding is set to {positional_embedding} and" \
+
+        assert seq_len is not None or positional_embedding == "none", (
+            f"Positional embedding is set to {positional_embedding} and"
             f" the sequence length was not specified."
-        
+        )
+
         if not seq_pool:
             seq_len += 1
-            self.class_emb = Parameter(torch.zeros(1, 1, self.embedding_dim),
-                                       requires_grad=True)
+            self.class_emb = Parameter(
+                torch.zeros(1, 1, self.embedding_dim), requires_grad=True
+            )
             self.num_tokens = 1
         else:
             self.attention_pool = Linear(self.embedding_dim, 1)
-        
-        if positional_embedding != 'none':
-            if positional_embedding == 'learnable':
+
+        if positional_embedding != "none":
+            if positional_embedding == "learnable":
                 seq_len += 1  # padding idx
-                self.positional_emb = Parameter(torch.zeros(1, seq_len, embedding_dim),
-                                                requires_grad=True)
+                self.positional_emb = Parameter(
+                    torch.zeros(1, seq_len, embedding_dim), requires_grad=True
+                )
                 init.trunc_normal_(self.positional_emb, std=0.2)
             else:
-                self.positional_emb = Parameter(self.sinusoidal_embedding(seq_len,
-                                                                          embedding_dim,
-                                                                          padding_idx=True),
-                                                requires_grad=False)
+                self.positional_emb = Parameter(
+                    self.sinusoidal_embedding(seq_len, embedding_dim, padding_idx=True),
+                    requires_grad=False,
+                )
         else:
             self.positional_emb = None
-        
+
         self.dropout = Dropout(p=dropout)
         dpr = [x.item() for x in torch.linspace(0, stochastic_depth, num_layers)]
-        self.blocks = ModuleList([
-            MaskedTransformerEncoderLayer(d_model=embedding_dim, nhead=num_heads,
-                                          dim_feedforward=dim_feedforward, dropout=dropout,
-                                          attention_dropout=attention_dropout, drop_path_rate=dpr[i])
-            for i in range(num_layers)])
+        self.blocks = ModuleList(
+            [
+                MaskedTransformerEncoderLayer(
+                    d_model=embedding_dim,
+                    nhead=num_heads,
+                    dim_feedforward=dim_feedforward,
+                    dropout=dropout,
+                    attention_dropout=attention_dropout,
+                    drop_path_rate=dpr[i],
+                )
+                for i in range(num_layers)
+            ]
+        )
         self.norm = LayerNorm(embedding_dim)
-        
+
         self.fc = Linear(embedding_dim, num_classes)
         self.apply(self.init_weight)
-    
+
     def forward(self, x, mask=None):
         if self.positional_emb is None and x.size(1) < self.seq_len:
-            x = F.pad(x, (0, 0, 0, self.n_channels - x.size(1)), mode='constant', value=0)
-        
+            x = F.pad(
+                x, (0, 0, 0, self.n_channels - x.size(1)), mode="constant", value=0
+            )
+
         if not self.seq_pool:
             cls_token = self.class_emb.expand(x.shape[0], -1, -1)
             x = torch.cat((cls_token, x), dim=1)
             if mask is not None:
-                mask = torch.cat([torch.ones(size=(mask.shape[0], 1), device=mask.device), mask.float()], dim=1)
-                mask = (mask > 0)
-        
+                mask = torch.cat(
+                    [
+                        torch.ones(size=(mask.shape[0], 1), device=mask.device),
+                        mask.float(),
+                    ],
+                    dim=1,
+                )
+                mask = mask > 0
+
         if self.positional_emb is not None:
             x += self.positional_emb
-        
+
         x = self.dropout(x)
-        
+
         for blk in self.blocks:
             x = blk(x, mask=mask)
         x = self.norm(x)
-        
+
         if self.seq_pool:
-            x = torch.matmul(F.softmax(self.attention_pool(x), dim=1).transpose(-1, -2), x).squeeze(-2)
+            x = torch.matmul(
+                F.softmax(self.attention_pool(x), dim=1).transpose(-1, -2), x
+            ).squeeze(-2)
         else:
             x = x[:, 0]
-        
+
         x = self.fc(x)
         return x
-    
+
     @staticmethod
     def init_weight(m):
         if isinstance(m, Linear):
-            init.trunc_normal_(m.weight, std=.02)
+            init.trunc_normal_(m.weight, std=0.02)
             if isinstance(m, Linear) and m.bias is not None:
                 init.constant_(m.bias, 0)
         elif isinstance(m, LayerNorm):
             init.constant_(m.bias, 0)
             init.constant_(m.weight, 1.0)
-    
+
     @staticmethod
     def sinusoidal_embedding(n_channels, dim, padding_idx=False):
-        pe = torch.FloatTensor([[p / (10000 ** (2 * (i // 2) / dim)) for i in range(dim)]
-                                for p in range(n_channels)])
+        pe = torch.FloatTensor(
+            [
+                [p / (10000 ** (2 * (i // 2) / dim)) for i in range(dim)]
+                for p in range(n_channels)
+            ]
+        )
         pe[:, 0::2] = torch.sin(pe[:, 0::2])
         pe[:, 1::2] = torch.cos(pe[:, 1::2])
         pe = pe.unsqueeze(0)
diff --git a/src/blades/models/cifar10/__init__.py b/src/blades/models/cifar10/__init__.py
index 659e5f4..aec8c2e 100644
--- a/src/blades/models/cifar10/__init__.py
+++ b/src/blades/models/cifar10/__init__.py
@@ -1,1 +1,5 @@
 from .cct import CCTNet
+
+__all__ = [
+    "CCTNet",
+]
diff --git a/src/blades/models/cifar10/cct.py b/src/blades/models/cifar10/cct.py
index 1afb825..8a83e5a 100644
--- a/src/blades/models/cifar10/cct.py
+++ b/src/blades/models/cifar10/cct.py
@@ -1,15 +1,15 @@
 import torch.nn as nn
 
-from .cctnets import cct_2_3x2_32
+from blades.models.backbones.cctnets import cct_2_3x2_32
 
 
 class CCTNet(nn.Module):
     def __init__(self):
         super(CCTNet, self).__init__()
-        self.mdoel = cct_2_3x2_32()
-    
+        self.model = cct_2_3x2_32()
+
     def forward(self, x):
-        return self.mdoel(x)
+        return self.model(x)
 
 
 def create_model():
diff --git a/src/blades/models/cifar10/cctnets/cct.py b/src/blades/models/cifar10/cctnets/cct.py
deleted file mode 100644
index f08cf70..0000000
--- a/src/blades/models/cifar10/cctnets/cct.py
+++ /dev/null
@@ -1,352 +0,0 @@
-import torch.nn as nn
-from torch.hub import load_state_dict_from_url
-
-from .utils.helpers import pe_check, fc_check
-from .utils.tokenizer import Tokenizer
-from .utils.transformers import TransformerClassifier
-
-try:
-    from timm.models.registry import register_model
-except ImportError:
-    from .registry import register_model
-
-model_urls = {
-    'cct_7_3x1_32':
-        'http://ix.cs.uoregon.edu/~alih/compact-transformers/checkpoints/pretrained/cct_7_3x1_32_cifar10_300epochs.pth',
-    'cct_7_3x1_32_sine':
-        'http://ix.cs.uoregon.edu/~alih/compact-transformers/checkpoints/pretrained/cct_7_3x1_32_sine_cifar10_5000epochs.pth',
-    'cct_7_3x1_32_c100':
-        'http://ix.cs.uoregon.edu/~alih/compact-transformers/checkpoints/pretrained/cct_7_3x1_32_cifar100_300epochs.pth',
-    'cct_7_3x1_32_sine_c100':
-        'http://ix.cs.uoregon.edu/~alih/compact-transformers/checkpoints/pretrained/cct_7_3x1_32_sine_cifar100_5000epochs.pth',
-    'cct_7_7x2_224_sine':
-        'http://ix.cs.uoregon.edu/~alih/compact-transformers/checkpoints/pretrained/cct_7_7x2_224_flowers102.pth',
-    'cct_14_7x2_224':
-        'http://ix.cs.uoregon.edu/~alih/compact-transformers/checkpoints/pretrained/cct_14_7x2_224_imagenet.pth',
-    'cct_14_7x2_384':
-        'http://ix.cs.uoregon.edu/~alih/compact-transformers/checkpoints/finetuned/cct_14_7x2_384_imagenet.pth',
-    'cct_14_7x2_384_fl':
-        'http://ix.cs.uoregon.edu/~alih/compact-transformers/checkpoints/finetuned/cct_14_7x2_384_flowers102.pth',
-}
-
-
-class CCT(nn.Module):
-    def __init__(self,
-                 img_size=224,
-                 embedding_dim=768,
-                 n_input_channels=3,
-                 n_conv_layers=1,
-                 kernel_size=7,
-                 stride=2,
-                 padding=3,
-                 pooling_kernel_size=3,
-                 pooling_stride=2,
-                 pooling_padding=1,
-                 dropout=0.,
-                 attention_dropout=0.1,
-                 stochastic_depth=0.1,
-                 num_layers=14,
-                 num_heads=6,
-                 mlp_ratio=4.0,
-                 num_classes=1000,
-                 positional_embedding='learnable',
-                 *args, **kwargs):
-        super(CCT, self).__init__()
-        
-        self.tokenizer = Tokenizer(n_input_channels=n_input_channels,
-                                   n_output_channels=embedding_dim,
-                                   kernel_size=kernel_size,
-                                   stride=stride,
-                                   padding=padding,
-                                   pooling_kernel_size=pooling_kernel_size,
-                                   pooling_stride=pooling_stride,
-                                   pooling_padding=pooling_padding,
-                                   max_pool=True,
-                                   activation=nn.ReLU,
-                                   n_conv_layers=n_conv_layers,
-                                   conv_bias=False)
-        
-        self.classifier = TransformerClassifier(
-            sequence_length=self.tokenizer.sequence_length(n_channels=n_input_channels,
-                                                           height=img_size,
-                                                           width=img_size),
-            embedding_dim=embedding_dim,
-            seq_pool=True,
-            dropout=dropout,
-            attention_dropout=attention_dropout,
-            stochastic_depth=stochastic_depth,
-            num_layers=num_layers,
-            num_heads=num_heads,
-            mlp_ratio=mlp_ratio,
-            num_classes=num_classes,
-            positional_embedding=positional_embedding
-        )
-    
-    def forward(self, x):
-        x = self.tokenizer(x)
-        return self.classifier(x)
-
-
-def _cct(arch, pretrained, progress,
-         num_layers, num_heads, mlp_ratio, embedding_dim,
-         kernel_size=3, stride=None, padding=None,
-         positional_embedding='learnable',
-         *args, **kwargs):
-    stride = stride if stride is not None else max(1, (kernel_size // 2) - 1)
-    padding = padding if padding is not None else max(1, (kernel_size // 2))
-    model = CCT(num_layers=num_layers,
-                num_heads=num_heads,
-                mlp_ratio=mlp_ratio,
-                embedding_dim=embedding_dim,
-                kernel_size=kernel_size,
-                stride=stride,
-                padding=padding,
-                *args, **kwargs)
-    
-    if pretrained:
-        if arch in model_urls:
-            state_dict = load_state_dict_from_url(model_urls[arch],
-                                                  progress=progress)
-            if positional_embedding == 'learnable':
-                state_dict = pe_check(model, state_dict)
-            elif positional_embedding == 'sine':
-                state_dict['classifier.positional_emb'] = model.state_dict()['classifier.positional_emb']
-            state_dict = fc_check(model, state_dict)
-            model.load_state_dict(state_dict)
-        else:
-            raise RuntimeError(f'Variant {arch} does not yet have pretrained weights.')
-    return model
-
-
-def cct_2(arch, pretrained, progress, *args, **kwargs):
-    return _cct(arch, pretrained, progress, num_layers=2, num_heads=2, mlp_ratio=1, embedding_dim=128,
-                *args, **kwargs)
-
-
-def cct_4(arch, pretrained, progress, *args, **kwargs):
-    return _cct(arch, pretrained, progress, num_layers=4, num_heads=2, mlp_ratio=1, embedding_dim=128,
-                *args, **kwargs)
-
-
-def cct_6(arch, pretrained, progress, *args, **kwargs):
-    return _cct(arch, pretrained, progress, num_layers=6, num_heads=4, mlp_ratio=2, embedding_dim=256,
-                *args, **kwargs)
-
-
-def cct_7(arch, pretrained, progress, *args, **kwargs):
-    return _cct(arch, pretrained, progress, num_layers=7, num_heads=4, mlp_ratio=2, embedding_dim=256,
-                *args, **kwargs)
-
-
-def cct_14(arch, pretrained, progress, *args, **kwargs):
-    return _cct(arch, pretrained, progress, num_layers=14, num_heads=6, mlp_ratio=3, embedding_dim=384,
-                *args, **kwargs)
-
-
-@register_model
-def cct_2_3x2_32(pretrained=False, progress=False,
-                 img_size=32, positional_embedding='learnable', num_classes=10,
-                 *args, **kwargs):
-    return cct_2('cct_2_3x2_32', pretrained, progress,
-                 kernel_size=3, n_conv_layers=2,
-                 img_size=img_size, positional_embedding=positional_embedding,
-                 num_classes=num_classes,
-                 *args, **kwargs)
-
-
-@register_model
-def cct_2_3x2_32_sine(pretrained=False, progress=False,
-                      img_size=32, positional_embedding='sine', num_classes=10,
-                      *args, **kwargs):
-    return cct_2('cct_2_3x2_32_sine', pretrained, progress,
-                 kernel_size=3, n_conv_layers=2,
-                 img_size=img_size, positional_embedding=positional_embedding,
-                 num_classes=num_classes,
-                 *args, **kwargs)
-
-
-@register_model
-def cct_4_3x2_32(pretrained=False, progress=False,
-                 img_size=32, positional_embedding='learnable', num_classes=10,
-                 *args, **kwargs):
-    return cct_4('cct_4_3x2_32', pretrained, progress,
-                 kernel_size=3, n_conv_layers=2,
-                 img_size=img_size, positional_embedding=positional_embedding,
-                 num_classes=num_classes,
-                 *args, **kwargs)
-
-
-@register_model
-def cct_4_3x2_32_sine(pretrained=False, progress=False,
-                      img_size=32, positional_embedding='sine', num_classes=10,
-                      *args, **kwargs):
-    return cct_4('cct_4_3x2_32_sine', pretrained, progress,
-                 kernel_size=3, n_conv_layers=2,
-                 img_size=img_size, positional_embedding=positional_embedding,
-                 num_classes=num_classes,
-                 *args, **kwargs)
-
-
-@register_model
-def cct_6_3x1_32(pretrained=False, progress=False,
-                 img_size=32, positional_embedding='learnable', num_classes=10,
-                 *args, **kwargs):
-    return cct_6('cct_6_3x1_32', pretrained, progress,
-                 kernel_size=3, n_conv_layers=1,
-                 img_size=img_size, positional_embedding=positional_embedding,
-                 num_classes=num_classes,
-                 *args, **kwargs)
-
-
-@register_model
-def cct_6_3x1_32_sine(pretrained=False, progress=False,
-                      img_size=32, positional_embedding='sine', num_classes=10,
-                      *args, **kwargs):
-    return cct_6('cct_6_3x1_32_sine', pretrained, progress,
-                 kernel_size=3, n_conv_layers=1,
-                 img_size=img_size, positional_embedding=positional_embedding,
-                 num_classes=num_classes,
-                 *args, **kwargs)
-
-
-@register_model
-def cct_6_3x2_32(pretrained=False, progress=False,
-                 img_size=32, positional_embedding='learnable', num_classes=10,
-                 *args, **kwargs):
-    return cct_6('cct_6_3x2_32', pretrained, progress,
-                 kernel_size=3, n_conv_layers=2,
-                 img_size=img_size, positional_embedding=positional_embedding,
-                 num_classes=num_classes,
-                 *args, **kwargs)
-
-
-@register_model
-def cct_6_3x2_32_sine(pretrained=False, progress=False,
-                      img_size=32, positional_embedding='sine', num_classes=10,
-                      *args, **kwargs):
-    return cct_6('cct_6_3x2_32_sine', pretrained, progress,
-                 kernel_size=3, n_conv_layers=2,
-                 img_size=img_size, positional_embedding=positional_embedding,
-                 num_classes=num_classes,
-                 *args, **kwargs)
-
-
-@register_model
-def cct_7_3x1_32(pretrained=False, progress=False,
-                 img_size=32, positional_embedding='learnable', num_classes=10,
-                 *args, **kwargs):
-    return cct_7('cct_7_3x1_32', pretrained, progress,
-                 kernel_size=3, n_conv_layers=1,
-                 img_size=img_size, positional_embedding=positional_embedding,
-                 num_classes=num_classes,
-                 *args, **kwargs)
-
-
-@register_model
-def cct_7_3x1_32_sine(pretrained=False, progress=False,
-                      img_size=32, positional_embedding='sine', num_classes=10,
-                      *args, **kwargs):
-    return cct_7('cct_7_3x1_32_sine', pretrained, progress,
-                 kernel_size=3, n_conv_layers=1,
-                 img_size=img_size, positional_embedding=positional_embedding,
-                 num_classes=num_classes,
-                 *args, **kwargs)
-
-
-@register_model
-def cct_7_3x1_32_c100(pretrained=False, progress=False,
-                      img_size=32, positional_embedding='learnable', num_classes=100,
-                      *args, **kwargs):
-    return cct_7('cct_7_3x1_32_c100', pretrained, progress,
-                 kernel_size=3, n_conv_layers=1,
-                 img_size=img_size, positional_embedding=positional_embedding,
-                 num_classes=num_classes,
-                 *args, **kwargs)
-
-
-@register_model
-def cct_7_3x1_32_sine_c100(pretrained=False, progress=False,
-                           img_size=32, positional_embedding='sine', num_classes=100,
-                           *args, **kwargs):
-    return cct_7('cct_7_3x1_32_sine_c100', pretrained, progress,
-                 kernel_size=3, n_conv_layers=1,
-                 img_size=img_size, positional_embedding=positional_embedding,
-                 num_classes=num_classes,
-                 *args, **kwargs)
-
-
-@register_model
-def cct_7_3x2_32(pretrained=False, progress=False,
-                 img_size=32, positional_embedding='learnable', num_classes=10,
-                 *args, **kwargs):
-    return cct_7('cct_7_3x2_32', pretrained, progress,
-                 kernel_size=3, n_conv_layers=2,
-                 img_size=img_size, positional_embedding=positional_embedding,
-                 num_classes=num_classes,
-                 *args, **kwargs)
-
-
-@register_model
-def cct_7_3x2_32_sine(pretrained=False, progress=False,
-                      img_size=32, positional_embedding='sine', num_classes=10,
-                      *args, **kwargs):
-    return cct_7('cct_7_3x2_32_sine', pretrained, progress,
-                 kernel_size=3, n_conv_layers=2,
-                 img_size=img_size, positional_embedding=positional_embedding,
-                 num_classes=num_classes,
-                 *args, **kwargs)
-
-
-@register_model
-def cct_7_7x2_224(pretrained=False, progress=False,
-                  img_size=224, positional_embedding='learnable', num_classes=102,
-                  *args, **kwargs):
-    return cct_7('cct_7_7x2_224', pretrained, progress,
-                 kernel_size=7, n_conv_layers=2,
-                 img_size=img_size, positional_embedding=positional_embedding,
-                 num_classes=num_classes,
-                 *args, **kwargs)
-
-
-@register_model
-def cct_7_7x2_224_sine(pretrained=False, progress=False,
-                       img_size=224, positional_embedding='sine', num_classes=102,
-                       *args, **kwargs):
-    return cct_7('cct_7_7x2_224_sine', pretrained, progress,
-                 kernel_size=7, n_conv_layers=2,
-                 img_size=img_size, positional_embedding=positional_embedding,
-                 num_classes=num_classes,
-                 *args, **kwargs)
-
-
-@register_model
-def cct_14_7x2_224(pretrained=False, progress=False,
-                   img_size=224, positional_embedding='learnable', num_classes=1000,
-                   *args, **kwargs):
-    return cct_14('cct_14_7x2_224', pretrained, progress,
-                  kernel_size=7, n_conv_layers=2,
-                  img_size=img_size, positional_embedding=positional_embedding,
-                  num_classes=num_classes,
-                  *args, **kwargs)
-
-
-@register_model
-def cct_14_7x2_384(pretrained=False, progress=False,
-                   img_size=384, positional_embedding='learnable', num_classes=1000,
-                   *args, **kwargs):
-    return cct_14('cct_14_7x2_384', pretrained, progress,
-                  kernel_size=7, n_conv_layers=2,
-                  img_size=img_size, positional_embedding=positional_embedding,
-                  num_classes=num_classes,
-                  *args, **kwargs)
-
-
-@register_model
-def cct_14_7x2_384_fl(pretrained=False, progress=False,
-                      img_size=384, positional_embedding='learnable', num_classes=102,
-                      *args, **kwargs):
-    return cct_14('cct_14_7x2_384_fl', pretrained, progress,
-                  kernel_size=7, n_conv_layers=2,
-                  img_size=img_size, positional_embedding=positional_embedding,
-                  num_classes=num_classes,
-                  *args, **kwargs)
diff --git a/src/blades/models/cifar10/cctnets/cvt.py b/src/blades/models/cifar10/cctnets/cvt.py
deleted file mode 100644
index d47962a..0000000
--- a/src/blades/models/cifar10/cctnets/cvt.py
+++ /dev/null
@@ -1,198 +0,0 @@
-import torch.nn as nn
-from torch.hub import load_state_dict_from_url
-
-from .utils.helpers import pe_check
-from .utils.tokenizer import Tokenizer
-from .utils.transformers import TransformerClassifier
-
-try:
-    from timm.models.registry import register_model
-except ImportError:
-    from .registry import register_model
-
-model_urls = {
-}
-
-
-class CVT(nn.Module):
-    def __init__(self,
-                 img_size=224,
-                 embedding_dim=768,
-                 n_input_channels=3,
-                 kernel_size=16,
-                 dropout=0.,
-                 attention_dropout=0.1,
-                 stochastic_depth=0.1,
-                 num_layers=14,
-                 num_heads=6,
-                 mlp_ratio=4.0,
-                 num_classes=1000,
-                 positional_embedding='learnable',
-                 *args, **kwargs):
-        super(CVT, self).__init__()
-        assert img_size % kernel_size == 0, f"Image size ({img_size}) has to be" \
-                                            f"divisible by patch size ({kernel_size})"
-        self.tokenizer = Tokenizer(n_input_channels=n_input_channels,
-                                   n_output_channels=embedding_dim,
-                                   kernel_size=kernel_size,
-                                   stride=kernel_size,
-                                   padding=0,
-                                   max_pool=False,
-                                   activation=None,
-                                   n_conv_layers=1,
-                                   conv_bias=True)
-        
-        self.classifier = TransformerClassifier(
-            sequence_length=self.tokenizer.sequence_length(n_channels=n_input_channels,
-                                                           height=img_size,
-                                                           width=img_size),
-            embedding_dim=embedding_dim,
-            seq_pool=True,
-            dropout=dropout,
-            attention_dropout=attention_dropout,
-            stochastic_depth=stochastic_depth,
-            num_layers=num_layers,
-            num_heads=num_heads,
-            mlp_ratio=mlp_ratio,
-            num_classes=num_classes,
-            positional_embedding=positional_embedding
-        )
-    
-    def forward(self, x):
-        x = self.tokenizer(x)
-        return self.classifier(x)
-
-
-def _cvt(arch, pretrained, progress,
-         num_layers, num_heads, mlp_ratio, embedding_dim,
-         kernel_size=4, positional_embedding='learnable',
-         *args, **kwargs):
-    model = CVT(num_layers=num_layers,
-                num_heads=num_heads,
-                mlp_ratio=mlp_ratio,
-                embedding_dim=embedding_dim,
-                kernel_size=kernel_size,
-                *args, **kwargs)
-    
-    if pretrained and arch in model_urls:
-        state_dict = load_state_dict_from_url(model_urls[arch],
-                                              progress=progress)
-        if positional_embedding == 'learnable':
-            state_dict = pe_check(model, state_dict)
-        elif positional_embedding == 'sine':
-            state_dict['classifier.positional_emb'] = model.state_dict()['classifier.positional_emb']
-        model.load_state_dict(state_dict)
-    return model
-
-
-def cvt_2(*args, **kwargs):
-    return _cvt(num_layers=2, num_heads=2, mlp_ratio=1, embedding_dim=128,
-                *args, **kwargs)
-
-
-def cvt_4(*args, **kwargs):
-    return _cvt(num_layers=4, num_heads=2, mlp_ratio=1, embedding_dim=128,
-                *args, **kwargs)
-
-
-def cvt_6(*args, **kwargs):
-    return _cvt(num_layers=6, num_heads=4, mlp_ratio=2, embedding_dim=256,
-                *args, **kwargs)
-
-
-def cvt_7(*args, **kwargs):
-    return _cvt(num_layers=7, num_heads=4, mlp_ratio=2, embedding_dim=256,
-                *args, **kwargs)
-
-
-def cvt_8(*args, **kwargs):
-    return _cvt(num_layers=8, num_heads=4, mlp_ratio=2, embedding_dim=256,
-                *args, **kwargs)
-
-
-@register_model
-def cvt_2_4_32(pretrained=False, progress=False,
-               img_size=32, positional_embedding='learnable', num_classes=10,
-               *args, **kwargs):
-    return cvt_2('cvt_2_4_32', pretrained, progress,
-                 kernel_size=4,
-                 img_size=img_size, positional_embedding=positional_embedding,
-                 num_classes=num_classes,
-                 *args, **kwargs)
-
-
-@register_model
-def cvt_2_4_32_sine(pretrained=False, progress=False,
-                    img_size=32, positional_embedding='sine', num_classes=10,
-                    *args, **kwargs):
-    return cvt_2('cvt_2_4_32_sine', pretrained, progress,
-                 kernel_size=4,
-                 img_size=img_size, positional_embedding=positional_embedding,
-                 num_classes=num_classes,
-                 *args, **kwargs)
-
-
-@register_model
-def cvt_4_4_32(pretrained=False, progress=False,
-               img_size=32, positional_embedding='learnable', num_classes=10,
-               *args, **kwargs):
-    return cvt_4('cvt_4_4_32', pretrained, progress,
-                 kernel_size=4,
-                 img_size=img_size, positional_embedding=positional_embedding,
-                 num_classes=num_classes,
-                 *args, **kwargs)
-
-
-@register_model
-def cvt_4_4_32_sine(pretrained=False, progress=False,
-                    img_size=32, positional_embedding='sine', num_classes=10,
-                    *args, **kwargs):
-    return cvt_4('cvt_4_4_32_sine', pretrained, progress,
-                 kernel_size=4,
-                 img_size=img_size, positional_embedding=positional_embedding,
-                 num_classes=num_classes,
-                 *args, **kwargs)
-
-
-@register_model
-def cvt_6_4_32(pretrained=False, progress=False,
-               img_size=32, positional_embedding='learnable', num_classes=10,
-               *args, **kwargs):
-    return cvt_6('cvt_6_4_32', pretrained, progress,
-                 kernel_size=4,
-                 img_size=img_size, positional_embedding=positional_embedding,
-                 num_classes=num_classes,
-                 *args, **kwargs)
-
-
-@register_model
-def cvt_6_4_32_sine(pretrained=False, progress=False,
-                    img_size=32, positional_embedding='sine', num_classes=10,
-                    *args, **kwargs):
-    return cvt_6('cvt_6_4_32_sine', pretrained, progress,
-                 kernel_size=4,
-                 img_size=img_size, positional_embedding=positional_embedding,
-                 num_classes=num_classes,
-                 *args, **kwargs)
-
-
-@register_model
-def cvt_7_4_32(pretrained=False, progress=False,
-               img_size=32, positional_embedding='learnable', num_classes=10,
-               *args, **kwargs):
-    return cvt_7('cvt_7_4_32', pretrained, progress,
-                 kernel_size=4,
-                 img_size=img_size, positional_embedding=positional_embedding,
-                 num_classes=num_classes,
-                 *args, **kwargs)
-
-
-@register_model
-def cvt_7_4_32_sine(pretrained=False, progress=False,
-                    img_size=32, positional_embedding='sine', num_classes=10,
-                    *args, **kwargs):
-    return cvt_7('cvt_7_4_32_sine', pretrained, progress,
-                 kernel_size=4,
-                 img_size=img_size, positional_embedding=positional_embedding,
-                 num_classes=num_classes,
-                 *args, **kwargs)
diff --git a/src/blades/models/cifar10/cctnets/text/__init__.py b/src/blades/models/cifar10/cctnets/text/__init__.py
deleted file mode 100644
index be499dd..0000000
--- a/src/blades/models/cifar10/cctnets/text/__init__.py
+++ /dev/null
@@ -1,3 +0,0 @@
-from .cct import *
-from .cvt import *
-from .vit import *
diff --git a/src/blades/models/cifar10/cctnets/text/cct.py b/src/blades/models/cifar10/cctnets/text/cct.py
deleted file mode 100644
index f29bc76..0000000
--- a/src/blades/models/cifar10/cctnets/text/cct.py
+++ /dev/null
@@ -1,86 +0,0 @@
-import torch.nn as nn
-
-from ..utils.embedder import Embedder
-from ..utils.tokenizer import TextTokenizer
-from ..utils.transformers import MaskedTransformerClassifier
-
-__all__ = [
-    'text_cct_2',
-    'text_cct_4',
-    'text_cct_6'
-]
-
-
-class TextCCT(nn.Module):
-    def __init__(self,
-                 seq_len=64,
-                 word_embedding_dim=300,
-                 embedding_dim=256,
-                 kernel_size=2,
-                 stride=1,
-                 padding=1,
-                 pooling_kernel_size=2,
-                 pooling_stride=2,
-                 pooling_padding=1,
-                 *args, **kwargs):
-        super(TextCCT, self).__init__()
-        
-        self.embedder = Embedder(word_embedding_dim=word_embedding_dim,
-                                 *args, **kwargs)
-        
-        self.tokenizer = TextTokenizer(n_input_channels=word_embedding_dim,
-                                       n_output_channels=embedding_dim,
-                                       kernel_size=kernel_size,
-                                       stride=stride,
-                                       padding=padding,
-                                       pooling_kernel_size=pooling_kernel_size,
-                                       pooling_stride=pooling_stride,
-                                       pooling_padding=pooling_padding,
-                                       max_pool=True,
-                                       activation=nn.ReLU)
-        
-        self.classifier = MaskedTransformerClassifier(
-            seq_len=self.tokenizer.seq_len(seq_len=seq_len, embed_dim=word_embedding_dim),
-            embedding_dim=embedding_dim,
-            seq_pool=True,
-            dropout=0.,
-            attention_dropout=0.1,
-            stochastic_depth=0.1,
-            *args, **kwargs)
-    
-    def forward(self, x, mask=None):
-        x, mask = self.embedder(x, mask=mask)
-        x, mask = self.tokenizer(x, mask=mask)
-        out = self.classifier(x, mask=mask)
-        return out
-
-
-def _text_cct(num_layers, num_heads, mlp_ratio, embedding_dim,
-              kernel_size=4, stride=None, padding=None,
-              *args, **kwargs):
-    stride = stride if stride is not None else max(1, (kernel_size // 2) - 1)
-    padding = padding if padding is not None else max(1, (kernel_size // 2))
-    
-    return TextCCT(num_layers=num_layers,
-                   num_heads=num_heads,
-                   mlp_ratio=mlp_ratio,
-                   embedding_dim=embedding_dim,
-                   kernel_size=kernel_size,
-                   stride=stride,
-                   padding=padding,
-                   *args, **kwargs)
-
-
-def text_cct_2(*args, **kwargs):
-    return _text_cct(num_layers=2, num_heads=2, mlp_ratio=1, embedding_dim=128,
-                     *args, **kwargs)
-
-
-def text_cct_4(*args, **kwargs):
-    return _text_cct(num_layers=4, num_heads=2, mlp_ratio=1, embedding_dim=128,
-                     *args, **kwargs)
-
-
-def text_cct_6(*args, **kwargs):
-    return _text_cct(num_layers=6, num_heads=4, mlp_ratio=2, embedding_dim=256,
-                     *args, **kwargs)
diff --git a/src/blades/models/cifar10/cctnets/text/cvt.py b/src/blades/models/cifar10/cctnets/text/cvt.py
deleted file mode 100644
index dad56d9..0000000
--- a/src/blades/models/cifar10/cctnets/text/cvt.py
+++ /dev/null
@@ -1,73 +0,0 @@
-import torch.nn as nn
-
-from ..utils.embedder import Embedder
-from ..utils.tokenizer import TextTokenizer
-from ..utils.transformers import MaskedTransformerClassifier
-
-__all__ = [
-    'text_cvt_2',
-    'text_cvt_4',
-    'text_cvt_6',
-]
-
-
-class TextCVT(nn.Module):
-    def __init__(self,
-                 seq_len=64,
-                 word_embedding_dim=300,
-                 embedding_dim=768,
-                 patch_size=2,
-                 *args, **kwargs):
-        super(TextCVT, self).__init__()
-        assert seq_len % patch_size == 0, f"sequence length ({seq_len}) has to be" \
-                                          f"divisible by patch size ({patch_size})"
-        self.embedder = Embedder(word_embedding_dim=word_embedding_dim,
-                                 *args, **kwargs)
-        
-        self.tokenizer = TextTokenizer(n_input_channels=word_embedding_dim,
-                                       n_output_channels=embedding_dim,
-                                       kernel_size=patch_size,
-                                       stride=patch_size,
-                                       padding=0,
-                                       max_pool=False,
-                                       activation=None)
-        
-        self.classifier = MaskedTransformerClassifier(
-            seq_len=self.tokenizer.seq_len(seq_len=seq_len, embed_dim=word_embedding_dim),
-            embedding_dim=embedding_dim,
-            seq_pool=True,
-            dropout=0.,
-            attention_dropout=0.1,
-            stochastic_depth=0.1,
-            *args, **kwargs)
-    
-    def forward(self, x, mask=None):
-        x, mask = self.embedder(x, mask=mask)
-        x, mask = self.tokenizer(x, mask=mask)
-        out = self.classifier(x, mask=mask)
-        return out
-
-
-def _text_cvt(num_layers, num_heads, mlp_ratio, embedding_dim,
-              patch_size=4, *args, **kwargs):
-    return TextCVT(num_layers=num_layers,
-                   num_heads=num_heads,
-                   mlp_ratio=mlp_ratio,
-                   embedding_dim=embedding_dim,
-                   patch_size=patch_size,
-                   *args, **kwargs)
-
-
-def text_cvt_2(*args, **kwargs):
-    return _text_cvt(num_layers=2, num_heads=2, mlp_ratio=1, embedding_dim=128,
-                     *args, **kwargs)
-
-
-def text_cvt_4(*args, **kwargs):
-    return _text_cvt(num_layers=4, num_heads=2, mlp_ratio=1, embedding_dim=128,
-                     *args, **kwargs)
-
-
-def text_cvt_6(*args, **kwargs):
-    return _text_cvt(num_layers=6, num_heads=4, mlp_ratio=2, embedding_dim=128,
-                     *args, **kwargs)
diff --git a/src/blades/models/cifar10/cctnets/text/transformer.py b/src/blades/models/cifar10/cctnets/text/transformer.py
deleted file mode 100644
index 49734a7..0000000
--- a/src/blades/models/cifar10/cctnets/text/transformer.py
+++ /dev/null
@@ -1,57 +0,0 @@
-import torch.nn as nn
-
-from ..utils.embedder import Embedder
-from ..utils.transformers import MaskedTransformerClassifier
-
-__all__ = [
-    'text_transformer_2',
-    'text_transformer_4',
-    'text_transformer_6',
-]
-
-
-class TextTransformerLite(nn.Module):
-    def __init__(self,
-                 seq_len=64,
-                 word_embedding_dim=300,
-                 *args, **kwargs):
-        super(TextTransformerLite, self).__init__()
-        self.embedder = Embedder(word_embedding_dim=word_embedding_dim,
-                                 *args, **kwargs)
-        
-        self.classifier = MaskedTransformerClassifier(
-            seq_len=seq_len,
-            embedding_dim=word_embedding_dim,
-            seq_pool=False,
-            dropout=0.,
-            attention_dropout=0.1,
-            stochastic_depth=0.1,
-            *args, **kwargs)
-    
-    def forward(self, x, mask=None):
-        x, mask = self.embedder(x, mask=mask)
-        out = self.classifier(x, mask=mask)
-        return out
-
-
-def _text_transformer(num_layers, num_heads, mlp_ratio,
-                      *args, **kwargs):
-    return TextTransformerLite(num_layers=num_layers,
-                               num_heads=num_heads,
-                               mlp_ratio=mlp_ratio,
-                               *args, **kwargs)
-
-
-def text_transformer_2(*args, **kwargs):
-    return _text_transformer(num_layers=2, num_heads=2, mlp_ratio=1,
-                             *args, **kwargs)
-
-
-def text_transformer_4(*args, **kwargs):
-    return _text_transformer(num_layers=4, num_heads=2, mlp_ratio=1,
-                             *args, **kwargs)
-
-
-def text_transformer_6(*args, **kwargs):
-    return _text_transformer(num_layers=6, num_heads=4, mlp_ratio=2,
-                             *args, **kwargs)
diff --git a/src/blades/models/cifar10/cctnets/text/vit.py b/src/blades/models/cifar10/cctnets/text/vit.py
deleted file mode 100644
index fe5ba38..0000000
--- a/src/blades/models/cifar10/cctnets/text/vit.py
+++ /dev/null
@@ -1,73 +0,0 @@
-import torch.nn as nn
-
-from ..utils.embedder import Embedder
-from ..utils.tokenizer import TextTokenizer
-from ..utils.transformers import MaskedTransformerClassifier
-
-__all__ = [
-    'text_vit_2',
-    'text_vit_4',
-    'text_vit_6',
-]
-
-
-class TextViTLite(nn.Module):
-    def __init__(self,
-                 seq_len=64,
-                 word_embedding_dim=300,
-                 embedding_dim=300,
-                 patch_size=2,
-                 *args, **kwargs):
-        super(TextViTLite, self).__init__()
-        assert seq_len % patch_size == 0, f"sequence length ({seq_len}) has to be" \
-                                          f"divisible by patch size ({patch_size})"
-        self.embedder = Embedder(word_embedding_dim=word_embedding_dim,
-                                 *args, **kwargs)
-        
-        self.tokenizer = TextTokenizer(n_input_channels=word_embedding_dim,
-                                       n_output_channels=embedding_dim,
-                                       kernel_size=patch_size,
-                                       stride=patch_size,
-                                       padding=0,
-                                       max_pool=False,
-                                       activation=None)
-        
-        self.classifier = MaskedTransformerClassifier(
-            seq_len=self.tokenizer.seq_len(seq_len=seq_len, embed_dim=word_embedding_dim),
-            embedding_dim=embedding_dim,
-            seq_pool=False,
-            dropout=0.,
-            attention_dropout=0.1,
-            stochastic_depth=0.1,
-            *args, **kwargs)
-    
-    def forward(self, x, mask=None):
-        x, mask = self.embedder(x, mask=mask)
-        x, mask = self.tokenizer(x, mask=mask)
-        out = self.classifier(x, mask=mask)
-        return out
-
-
-def _text_vit(num_layers, num_heads, mlp_ratio, embedding_dim,
-              patch_size=4, *args, **kwargs):
-    return TextViTLite(num_layers=num_layers,
-                       num_heads=num_heads,
-                       mlp_ratio=mlp_ratio,
-                       embedding_dim=embedding_dim,
-                       patch_size=patch_size,
-                       *args, **kwargs)
-
-
-def text_vit_2(*args, **kwargs):
-    return _text_vit(num_layers=2, num_heads=2, mlp_ratio=1, embedding_dim=128,
-                     *args, **kwargs)
-
-
-def text_vit_4(*args, **kwargs):
-    return _text_vit(num_layers=4, num_heads=2, mlp_ratio=1, embedding_dim=128,
-                     *args, **kwargs)
-
-
-def text_vit_6(*args, **kwargs):
-    return _text_vit(num_layers=6, num_heads=4, mlp_ratio=2, embedding_dim=256,
-                     *args, **kwargs)
diff --git a/src/blades/models/cifar10/cctnets/utils/embedder.py b/src/blades/models/cifar10/cctnets/utils/embedder.py
deleted file mode 100644
index a930dd9..0000000
--- a/src/blades/models/cifar10/cctnets/utils/embedder.py
+++ /dev/null
@@ -1,37 +0,0 @@
-import torch.nn as nn
-
-
-class Embedder(nn.Module):
-    def __init__(self,
-                 word_embedding_dim=300,
-                 vocab_size=100000,
-                 padding_idx=1,
-                 pretrained_weight=None,
-                 embed_freeze=False,
-                 *args, **kwargs):
-        super(Embedder, self).__init__()
-        self.embeddings = nn.Embedding.from_pretrained(pretrained_weight, freeze=embed_freeze) \
-            if pretrained_weight is not None else \
-            nn.Embedding(vocab_size, word_embedding_dim, padding_idx=padding_idx)
-        self.embeddings.weight.requires_grad = not embed_freeze
-    
-    def forward_mask(self, mask):
-        bsz, seq_len = mask.shape
-        new_mask = mask.view(bsz, seq_len, 1)
-        new_mask = new_mask.sum(-1)
-        new_mask = (new_mask > 0)
-        return new_mask
-    
-    def forward(self, x, mask=None):
-        embed = self.embeddings(x)
-        embed = embed if mask is None else embed * self.forward_mask(mask).unsqueeze(-1).float()
-        return embed, mask
-    
-    @staticmethod
-    def init_weight(m):
-        if isinstance(m, nn.Linear):
-            nn.init.trunc_normal_(m.weight, std=.02)
-            if isinstance(m, nn.Linear) and m.bias is not None:
-                nn.init.constant_(m.bias, 0)
-        else:
-            nn.init.normal_(m.weight)
diff --git a/src/blades/models/cifar10/cctnets/utils/tokenizer.py b/src/blades/models/cifar10/cctnets/utils/tokenizer.py
deleted file mode 100644
index 3e1bd5b..0000000
--- a/src/blades/models/cifar10/cctnets/utils/tokenizer.py
+++ /dev/null
@@ -1,111 +0,0 @@
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-
-class Tokenizer(nn.Module):
-    def __init__(self,
-                 kernel_size, stride, padding,
-                 pooling_kernel_size=3, pooling_stride=2, pooling_padding=1,
-                 n_conv_layers=1,
-                 n_input_channels=3,
-                 n_output_channels=64,
-                 in_planes=64,
-                 activation=None,
-                 max_pool=True,
-                 conv_bias=False):
-        super(Tokenizer, self).__init__()
-        
-        n_filter_list = [n_input_channels] + \
-                        [in_planes for _ in range(n_conv_layers - 1)] + \
-                        [n_output_channels]
-        
-        self.conv_layers = nn.Sequential(
-            *[nn.Sequential(
-                nn.Conv2d(n_filter_list[i], n_filter_list[i + 1],
-                          kernel_size=(kernel_size, kernel_size),
-                          stride=(stride, stride),
-                          padding=(padding, padding), bias=conv_bias),
-                nn.Identity() if activation is None else activation(),
-                nn.MaxPool2d(kernel_size=pooling_kernel_size,
-                             stride=pooling_stride,
-                             padding=pooling_padding) if max_pool else nn.Identity()
-            )
-                for i in range(n_conv_layers)
-            ])
-        
-        self.flattener = nn.Flatten(2, 3)
-        self.apply(self.init_weight)
-    
-    def sequence_length(self, n_channels=3, height=224, width=224):
-        return self.forward(torch.zeros((1, n_channels, height, width))).shape[1]
-    
-    def forward(self, x):
-        return self.flattener(self.conv_layers(x)).transpose(-2, -1)
-    
-    @staticmethod
-    def init_weight(m):
-        if isinstance(m, nn.Conv2d):
-            nn.init.kaiming_normal_(m.weight)
-
-
-class TextTokenizer(nn.Module):
-    def __init__(self,
-                 kernel_size, stride, padding,
-                 pooling_kernel_size=3, pooling_stride=2, pooling_padding=1,
-                 embedding_dim=300,
-                 n_output_channels=128,
-                 activation=None,
-                 max_pool=True,
-                 *args, **kwargs):
-        super(TextTokenizer, self).__init__()
-        
-        self.max_pool = max_pool
-        self.conv_layers = nn.Sequential(
-            nn.Conv2d(1, n_output_channels,
-                      kernel_size=(kernel_size, embedding_dim),
-                      stride=(stride, 1),
-                      padding=(padding, 0), bias=False),
-            nn.Identity() if activation is None else activation(),
-            nn.MaxPool2d(
-                kernel_size=(pooling_kernel_size, 1),
-                stride=(pooling_stride, 1),
-                padding=(pooling_padding, 0)
-            ) if max_pool else nn.Identity()
-        )
-        
-        self.apply(self.init_weight)
-    
-    def seq_len(self, seq_len=32, embed_dim=300):
-        return self.forward(torch.zeros((1, seq_len, embed_dim)))[0].shape[1]
-    
-    def forward_mask(self, mask):
-        new_mask = mask.unsqueeze(1).float()
-        cnn_weight = torch.ones(
-            (1, 1, self.conv_layers[0].kernel_size[0]),
-            device=mask.device,
-            dtype=torch.float)
-        new_mask = F.conv1d(
-            new_mask, cnn_weight, None,
-            self.conv_layers[0].stride[0], self.conv_layers[0].padding[0], 1, 1)
-        if self.max_pool:
-            new_mask = F.max_pool1d(
-                new_mask, self.conv_layers[2].kernel_size[0],
-                self.conv_layers[2].stride[0], self.conv_layers[2].padding[0], 1, False, False)
-        new_mask = new_mask.squeeze(1)
-        new_mask = (new_mask > 0)
-        return new_mask
-    
-    def forward(self, x, mask=None):
-        x = x.unsqueeze(1)
-        x = self.conv_layers(x)
-        x = x.transpose(1, 3).squeeze(1)
-        if mask is not None:
-            mask = self.forward_mask(mask).unsqueeze(-1).float()
-            x = x * mask
-        return x, mask
-    
-    @staticmethod
-    def init_weight(m):
-        if isinstance(m, nn.Conv2d):
-            nn.init.kaiming_normal_(m.weight)
diff --git a/src/blades/models/cifar10/cctnets/vit.py b/src/blades/models/cifar10/cctnets/vit.py
deleted file mode 100644
index 5e8ac22..0000000
--- a/src/blades/models/cifar10/cctnets/vit.py
+++ /dev/null
@@ -1,194 +0,0 @@
-import torch.nn as nn
-from torch.hub import load_state_dict_from_url
-
-from .utils.helpers import pe_check
-from .utils.tokenizer import Tokenizer
-from .utils.transformers import TransformerClassifier
-
-try:
-    from timm.models.registry import register_model
-except ImportError:
-    from .registry import register_model
-
-model_urls = {
-}
-
-
-class ViTLite(nn.Module):
-    def __init__(self,
-                 img_size=224,
-                 embedding_dim=768,
-                 n_input_channels=3,
-                 kernel_size=16,
-                 dropout=0.,
-                 attention_dropout=0.1,
-                 stochastic_depth=0.1,
-                 num_layers=14,
-                 num_heads=6,
-                 mlp_ratio=4.0,
-                 num_classes=1000,
-                 positional_embedding='learnable',
-                 *args, **kwargs):
-        super(ViTLite, self).__init__()
-        assert img_size % kernel_size == 0, f"Image size ({img_size}) has to be" \
-                                            f"divisible by patch size ({kernel_size})"
-        self.tokenizer = Tokenizer(n_input_channels=n_input_channels,
-                                   n_output_channels=embedding_dim,
-                                   kernel_size=kernel_size,
-                                   stride=kernel_size,
-                                   padding=0,
-                                   max_pool=False,
-                                   activation=None,
-                                   n_conv_layers=1,
-                                   conv_bias=True)
-        
-        self.classifier = TransformerClassifier(
-            sequence_length=self.tokenizer.sequence_length(n_channels=n_input_channels,
-                                                           height=img_size,
-                                                           width=img_size),
-            embedding_dim=embedding_dim,
-            seq_pool=False,
-            dropout=dropout,
-            attention_dropout=attention_dropout,
-            stochastic_depth=stochastic_depth,
-            num_layers=num_layers,
-            num_heads=num_heads,
-            mlp_ratio=mlp_ratio,
-            num_classes=num_classes,
-            positional_embedding=positional_embedding
-        )
-    
-    def forward(self, x):
-        x = self.tokenizer(x)
-        return self.classifier(x)
-
-
-def _vit_lite(arch, pretrained, progress,
-              num_layers, num_heads, mlp_ratio, embedding_dim,
-              kernel_size=4, *args, **kwargs):
-    model = ViTLite(num_layers=num_layers,
-                    num_heads=num_heads,
-                    mlp_ratio=mlp_ratio,
-                    embedding_dim=embedding_dim,
-                    kernel_size=kernel_size,
-                    positional_embedding='learnable',
-                    *args, **kwargs)
-    
-    if pretrained and arch in model_urls:
-        state_dict = load_state_dict_from_url(model_urls[arch],
-                                              progress=progress)
-        if positional_embedding == 'learnable':
-            state_dict = pe_check(model, state_dict)
-        elif positional_embedding == 'sine':
-            state_dict['classifier.positional_emb'] = model.state_dict()['classifier.positional_emb']
-        state_dict = pe_check(model, state_dict)
-        model.load_state_dict(state_dict)
-    return model
-
-
-def vit_2(*args, **kwargs):
-    return _vit_lite(num_layers=2, num_heads=2, mlp_ratio=1, embedding_dim=128,
-                     *args, **kwargs)
-
-
-def vit_4(*args, **kwargs):
-    return _vit_lite(num_layers=4, num_heads=2, mlp_ratio=1, embedding_dim=128,
-                     *args, **kwargs)
-
-
-def vit_6(*args, **kwargs):
-    return _vit_lite(num_layers=6, num_heads=4, mlp_ratio=2, embedding_dim=256,
-                     *args, **kwargs)
-
-
-def vit_7(*args, **kwargs):
-    return _vit_lite(num_layers=7, num_heads=4, mlp_ratio=2, embedding_dim=256,
-                     *args, **kwargs)
-
-
-@register_model
-def vit_2_4_32(pretrained=False, progress=False,
-               img_size=32, positional_embedding='learnable', num_classes=10,
-               *args, **kwargs):
-    return vit_2('vit_2_4_32', pretrained, progress,
-                 kernel_size=4,
-                 img_size=img_size, positional_embedding=positional_embedding,
-                 num_classes=num_classes,
-                 *args, **kwargs)
-
-
-@register_model
-def vit_2_4_32_sine(pretrained=False, progress=False,
-                    img_size=32, positional_embedding='sine', num_classes=10,
-                    *args, **kwargs):
-    return vit_2('vit_2_4_32_sine', pretrained, progress,
-                 kernel_size=4,
-                 img_size=img_size, positional_embedding=positional_embedding,
-                 num_classes=num_classes,
-                 *args, **kwargs)
-
-
-@register_model
-def vit_4_4_32(pretrained=False, progress=False,
-               img_size=32, positional_embedding='learnable', num_classes=10,
-               *args, **kwargs):
-    return vit_4('vit_4_4_32', pretrained, progress,
-                 kernel_size=4,
-                 img_size=img_size, positional_embedding=positional_embedding,
-                 num_classes=num_classes,
-                 *args, **kwargs)
-
-
-@register_model
-def vit_4_4_32_sine(pretrained=False, progress=False,
-                    img_size=32, positional_embedding='sine', num_classes=10,
-                    *args, **kwargs):
-    return vit_4('vit_4_4_32_sine', pretrained, progress,
-                 kernel_size=4,
-                 img_size=img_size, positional_embedding=positional_embedding,
-                 num_classes=num_classes,
-                 *args, **kwargs)
-
-
-@register_model
-def vit_6_4_32(pretrained=False, progress=False,
-               img_size=32, positional_embedding='learnable', num_classes=10,
-               *args, **kwargs):
-    return vit_6('vit_6_4_32', pretrained, progress,
-                 kernel_size=4,
-                 img_size=img_size, positional_embedding=positional_embedding,
-                 num_classes=num_classes,
-                 *args, **kwargs)
-
-
-@register_model
-def vit_6_4_32_sine(pretrained=False, progress=False,
-                    img_size=32, positional_embedding='sine', num_classes=10,
-                    *args, **kwargs):
-    return vit_6('vit_6_4_32_sine', pretrained, progress,
-                 kernel_size=4,
-                 img_size=img_size, positional_embedding=positional_embedding,
-                 num_classes=num_classes,
-                 *args, **kwargs)
-
-
-@register_model
-def vit_7_4_32(pretrained=False, progress=False,
-               img_size=32, positional_embedding='learnable', num_classes=10,
-               *args, **kwargs):
-    return vit_7('vit_7_4_32', pretrained, progress,
-                 kernel_size=4,
-                 img_size=img_size, positional_embedding=positional_embedding,
-                 num_classes=num_classes,
-                 *args, **kwargs)
-
-
-@register_model
-def vit_7_4_32_sine(pretrained=False, progress=False,
-                    img_size=32, positional_embedding='sine', num_classes=10,
-                    *args, **kwargs):
-    return vit_7('vit_7_4_32_sine', pretrained, progress,
-                 kernel_size=4,
-                 img_size=img_size, positional_embedding=positional_embedding,
-                 num_classes=num_classes,
-                 *args, **kwargs)
diff --git a/src/blades/models/cifar10/get_cifar10.py b/src/blades/models/cifar10/get_cifar10.py
deleted file mode 100644
index f887d91..0000000
--- a/src/blades/models/cifar10/get_cifar10.py
+++ /dev/null
@@ -1,89 +0,0 @@
-import argparse
-import os
-import pickle
-
-import numpy as np
-import torchvision
-from sklearn.utils import shuffle
-
-
-def generate_datasets(iid=False, alpha=1.0, num_clients=100):
-    train_set = torchvision.datasets.CIFAR10(train=True, download=True, root="./data")
-    test_set = torchvision.datasets.CIFAR10(train=False, download=True, root="./data")
-    x_test, y_test = test_set.data, np.array(test_set.targets)
-    x_train, y_train = train_set.data, np.array(train_set.targets)
-    
-    # (x_train, y_train), (x_test, y_test) = cifar10.load_data()
-    x_train = x_train.astype('float32') / 255.0
-    x_train = np.transpose(x_train, (0, 3, 1, 2))
-    x_test = x_test.astype('float32') / 255.0
-    x_test = np.transpose(x_test, (0, 3, 1, 2))
-    
-    np.random.seed(1234)
-    x_train, y_train = shuffle(x_train, y_train)
-    x_test, y_test = shuffle(x_test, y_test)
-    
-    train_user_ids = [str(id) for id in range(num_clients)]
-    x_test_splits = np.split(x_test, num_clients)
-    y_test_splits = np.split(y_test, num_clients)
-    
-    if iid:
-        x_train_splits = np.split(x_train, num_clients)
-        y_train_splits = np.split(y_train, num_clients)
-    else:
-        print('generating non-iid data')
-        min_size = 0
-        K = 10
-        N = y_train.shape[0]
-        client_dataidx_map = {}
-        
-        while min_size < 10:
-            proportion_list = []
-            idx_batch = [[] for _ in range(num_clients)]
-            # for each class in the dataset
-            for k in range(K):
-                idx_k = np.where(y_train == k)[0]
-                np.random.shuffle(idx_k)
-                proportions = np.random.dirichlet(np.repeat(alpha, num_clients))
-                
-                ## Balance
-                proportions = np.array([p * (len(idx_j) < N / num_clients) for p, idx_j in zip(proportions, idx_batch)])
-                proportions = proportions / proportions.sum()
-                proportions = (np.cumsum(proportions) * len(idx_k)).astype(int)[:-1]
-                print(proportions)
-                idx_batch = [idx_j + idx.tolist() for idx_j, idx in zip(idx_batch, np.split(idx_k, proportions))]
-                min_size = min([len(idx_j) for idx_j in idx_batch])
-                proportion_list.append(proportions)
-        x_train_splits, y_train_splits = [], []
-        for j in range(num_clients):
-            np.random.shuffle(idx_batch[j])
-            client_dataidx_map[j] = idx_batch[j]
-            x_train_splits.append(x_train[idx_batch[j], :])
-            y_train_splits.append(y_train[idx_batch[j], :])
-    
-    test_dataset = {}
-    train_dataset = {}
-    for id, index in zip(train_user_ids, range(num_clients)):
-        train_dataset[id] = {'x': x_train_splits[index], 'y': y_train_splits[index].flatten()}
-        test_dataset[id] = {'x': x_test_splits[index], 'y': y_test_splits[index].flatten()}
-    
-    #     os.system('rm -rf ..')
-    #     os.system('mkdir -p ../data')
-    with open(os.path.join('data', 'data_cache' + ("_alpha" + str(args.alpha) if not args.iid else "") + '.obj'),
-              'wb') as f:
-        pickle.dump(train_user_ids, f)
-        pickle.dump(train_dataset, f)
-        pickle.dump(train_user_ids, f)
-        pickle.dump(test_dataset, f)
-
-
-if __name__ == '__main__':
-    parser = argparse.ArgumentParser()
-    flag_parser = parser.add_mutually_exclusive_group(required=False)
-    flag_parser.add_argument('--iid', dest='iid', action='store_true')
-    flag_parser.add_argument('--noniid', dest='iid', action='store_false')
-    parser.add_argument('--alpha', type=float, default=0.1, required=False)
-    parser.add_argument('--num_clients', type=int, default=20, required=False)
-    parser.set_defaults(iid=True)
-    args = parser.parse_args()
-    generate_datasets(args.iid, alpha=args.alpha, num_clients=args.num_clients)
diff --git a/src/blades/models/cifar100/__init__.py b/src/blades/models/cifar100/__init__.py
new file mode 100644
index 0000000..aec8c2e
--- /dev/null
+++ b/src/blades/models/cifar100/__init__.py
@@ -0,0 +1,5 @@
+from .cct import CCTNet
+
+__all__ = [
+    "CCTNet",
+]
diff --git a/src/blades/models/cifar100/cct.py b/src/blades/models/cifar100/cct.py
new file mode 100644
index 0000000..41bc29d
--- /dev/null
+++ b/src/blades/models/cifar100/cct.py
@@ -0,0 +1,16 @@
+import torch.nn as nn
+
+from blades.models.backbones.cctnets import cct_7_3x1_32_c100
+
+
+class CCTNet(nn.Module):
+    def __init__(self):
+        super(CCTNet, self).__init__()
+        self.model = cct_7_3x1_32_c100()
+
+    def forward(self, x):
+        return self.model(x)
+
+
+def create_model():
+    return CCTNet(), nn.modules.loss.CrossEntropyLoss()
diff --git a/src/blades/models/mnist/__init__.py b/src/blades/models/mnist/__init__.py
index 8b055f5..cbe72a5 100644
--- a/src/blades/models/mnist/__init__.py
+++ b/src/blades/models/mnist/__init__.py
@@ -1,1 +1,5 @@
-from .dnn import MLP
+from .mlp import MLP
+
+__all__ = [
+    "MLP",
+]
diff --git a/src/blades/models/mnist/dnn.py b/src/blades/models/mnist/mlp.py
similarity index 99%
rename from src/blades/models/mnist/dnn.py
rename to src/blades/models/mnist/mlp.py
index 2121e6b..27119c3 100644
--- a/src/blades/models/mnist/dnn.py
+++ b/src/blades/models/mnist/mlp.py
@@ -1,7 +1,6 @@
 import torch.nn.functional as F
 from torch import nn
 
-
 # class MLP(nn.Module):
 #     def __init__(self):
 #         super(MLP, self).__init__()
@@ -9,7 +8,7 @@
 #         self.layer1 = nn.Linear(28 * 28, 64)
 #         self.layer2 = nn.Linear(64, 128)
 #         self.layer3 = nn.Linear(128, 10)
-    
+
 #     def forward(self, x):
 #         x = self.flatten(x)
 #         x = F.relu(self.layer1(x))
diff --git a/src/blades/models/utils/constants.py b/src/blades/models/utils/constants.py
deleted file mode 100644
index baef471..0000000
--- a/src/blades/models/utils/constants.py
+++ /dev/null
@@ -1,2 +0,0 @@
-DATASETS = ['sent140', 'femnist', 'shakespeare', 'celeba', 'synthetic']
-SEED_FILES = {'sampling': 'sampling_seed.txt', 'split': 'split_seed.txt'}
diff --git a/src/blades/models/utils/download_util.py b/src/blades/models/utils/download_util.py
deleted file mode 100644
index 282abf1..0000000
--- a/src/blades/models/utils/download_util.py
+++ /dev/null
@@ -1,45 +0,0 @@
-# taken from this StackOverflow answer: https://stackoverflow.com/a/39225039
-import os
-
-import requests
-
-
-def download_file_from_google_drive(id, destination):
-    URL = "https://docs.google.com/uc?export=download"
-    
-    session = requests.Session()
-    
-    response = session.get(URL, params={'_id': id}, stream=True)
-    token = get_confirm_token(response)
-    
-    if token:
-        params = {'_id': id, 'confirm': token}
-        response = session.get(URL, params=params, stream=True)
-    
-    save_response_content(response, destination)
-
-
-def get_confirm_token(response):
-    for key, value in response.cookies.items():
-        if key.startswith('download_warning'):
-            return value
-    
-    return None
-
-
-def save_response_content(response, destination):
-    CHUNK_SIZE = 32768
-    
-    with open(destination, "wb") as f:
-        for chunk in response.iter_content(CHUNK_SIZE):
-            if chunk:  # filter out keep-alive new chunks
-                f.write(chunk)
-
-
-if __name__ == "__main__":
-    file_id = "1rdRFbKeT9woS48Fmmo2mgJWDWSexhGeS"
-    destination = './femnist.zip'
-    download_file_from_google_drive(file_id, destination)
-    os.system('mkdir -p ../data')
-    os.system('unzip -o ' + destination + " -d ../data")
-    os.system('rm ' + destination)
diff --git a/src/blades/models/utils/preprocess.sh b/src/blades/models/utils/preprocess.sh
deleted file mode 100755
index 9a4d0eb..0000000
--- a/src/blades/models/utils/preprocess.sh
+++ /dev/null
@@ -1,255 +0,0 @@
-#!/usr/bin/env bash
-#
-# script to preprocess data
-
-# --------------------
-# parse arguments
-
-NAME="sent140" # name of the dataset, equivalent to directory name
-SAMPLE="na" # -s tag, iid or niid
-IUSER="" # --iu tag, # of users if iid sampling
-SFRAC="" # --sf tag, fraction of data to sample
-MINSAMPLES="na" # -k tag, minimum allowable # of samples per user
-TRAIN="na" # -t tag, user or sample
-TFRAC="" # --tf tag, fraction of data in training set
-SAMPLING_SEED="" # --smplseed, seed specified for sampling of data
-SPLIT_SEED="" # --spltseed, seed specified for train/test data split
-NO_CHECKSUM="" # --nochecksum, disable creation of MD5 checksum file after data gen
-VERIFICATION_FILE="" # --verify <fname>, check if JSON files' MD5 matches given digest
-
-META_DIR='meta'
-CHECKSUM_FNAME="${META_DIR}/dir-checksum.md5"
-
-while [[ $# -gt 0 ]]
-do
-key="$1"
-
-# TODO: Use getopts instead of creating cases!
-case $key in
-    --name)
-    NAME="$2"
-    shift # past argument
-    if [ ${SAMPLE:0:1} = "-" ]; then
-        NAME="sent140"
-    else
-        shift # past value
-    fi
-    ;;
-    -s)
-    SAMPLE="$2"
-    shift # past argument
-    if [ ${SAMPLE:0:1} = "-" ]; then
-        SAMPLE=""
-    else
-        shift # past value
-    fi
-    ;;
-    --iu)
-    IUSER="$2"
-    shift # past argument
-    if [ ${IUSER:0:1} = "-" ]; then
-        IUSER=""
-    else
-        shift # past value
-    fi
-    ;;
-    --sf)
-    SFRAC="$2"
-    shift # past argument
-    if [ ${SFRAC:0:1} = "-" ]; then
-        SFRAC=""
-    else
-        shift # past value
-    fi
-    ;;
-    -k)
-    MINSAMPLES="$2"
-    shift # past argument
-    if [ ${MINSAMPLES:0:1} = "-" ]; then
-        MINSAMPLES=""
-    else
-        shift # past value
-    fi
-    ;;
-    -t)
-    TRAIN="$2"
-    shift # past argument
-    if [ -z "$TRAIN" ] || [ ${TRAIN:0:1} = "-" ]; then
-        TRAIN=""
-    else
-        shift # past value
-    fi
-    ;;
-    --tf)
-    TFRAC="$2"
-    shift # past argument
-    if [ ${TFRAC:0:1} = "-" ]; then
-        TFRAC=""
-    else
-        shift # past value
-    fi
-    ;;
-    --smplseed)
-    SAMPLING_SEED="$2"
-    shift # past argument
-    ;;
-    --spltseed)
-    SPLIT_SEED="$2"
-    shift # past argument
-    ;;
-    --nochecksum)
-    NO_CHECKSUM="true"
-    shift # past argument
-    ;;
-    --verify)
-    VERIFICATION_FILE="$2"
-    shift # past argument
-    ;;
-    *)    # unknown option
-    shift # past argument
-    ;;
-esac
-done
-
-# --------------------
-# check if _running in verification mode
-
-if [ -n "${VERIFICATION_FILE}" ]; then
-    pushd ../$NAME >/dev/null 2>/dev/null
-        TMP_FILE=`mktemp /tmp/dir-checksum.XXXXXXX`
-        find 'data/' -type f -name '*.json' -exec md5sum {} + | sort -k 2 > ${TMP_FILE}
-        DIFF=`diff --brief ${VERIFICATION_FILE} ${TMP_FILE}`
-        if [ $? -ne 0 ]; then
-            echo "${DIFF}"
-            diff ${TMP_FILE} ${VERIFICATION_FILE} 
-            echo "Differing checksums found - please verify"
-        else
-            echo "Matching JSON files and checksums found!"
-        fi
-    popd >/dev/null 2>/dev/null
-    exit 0
-fi
-
-# --------------------
-# preprocess data
-
-CONT_SCRIPT=true
-cd ../$NAME
-
-# setup meta directory if doesn't exist
-if [ ! -d ${META_DIR} ]; then
-    mkdir -p ${META_DIR}
-fi
-META_DIR=`realpath ${META_DIR}`
-
-# download data and convert to .json format
-
-if [ ! -d "data/all_data" ]; then
-    cd preprocess
-    ./data_to_json.sh
-    cd ..
-fi
-
-NAMETAG="--name $NAME"
-
-# sample data
-IUSERTAG=""
-if [ ! -z $IUSER ]; then
-    IUSERTAG="--u $IUSER"
-fi
-SFRACTAG=""
-if [ ! -z $SFRAC ]; then
-    SFRACTAG="--fraction $SFRAC"
-fi
-
-if [ "$CONT_SCRIPT" = true ] && [ ! $SAMPLE = "na" ]; then
-    if [ -d "data/sampled_data" ] && [ "$(ls -A data/sampled_data)" ]; then
-        CONT_SCRIPT=false
-    else
-        if [ ! -d "data/sampled_data" ]; then
-            mkdir data/sampled_data
-        fi
-
-        cd ../utils
-
-        # Defaults to -1 if not specified, causes script to randomly generate seed
-        SEED_ARGUMENT="${SAMPLING_SEED:--1}" 
-
-        if [ $SAMPLE = "iid" ]; then
-            LEAF_DATA_META_DIR=${META_DIR} python3 sample.py $NAMETAG --iid $IUSERTAG $SFRACTAG --seed ${SEED_ARGUMENT}
-        else
-            LEAF_DATA_META_DIR=${META_DIR} python3 sample.py $NAMETAG $SFRACTAG --seed ${SEED_ARGUMENT}
-        fi
-
-        cd ../$NAME
-    fi
-fi
-
-# remove users with less then given number of samples
-if [ "$CONT_SCRIPT" = true ] && [ ! $MINSAMPLES = "na" ]; then
-    if [ -d "data/rem_user_data" ] && [ "$(ls -A data/rem_user_data)" ]; then
-        CONT_SCRIPT=false
-    else
-        if [ ! -d "data/rem_user_data" ]; then
-            mkdir data/rem_user_data
-        fi
-
-        cd ../utils
-
-        if [ -z $MINSAMPLES ]; then
-            python3 remove_users.py $NAMETAG
-        else
-            python3 remove_users.py $NAMETAG --min_samples $MINSAMPLES
-        fi
-
-        cd ../$NAME
-    fi
-fi
-
-# create train-test split
-TFRACTAG=""
-if [ ! -z $TFRAC ]; then
-    TFRACTAG="--frac $TFRAC"
-fi
-
-if [ "$CONT_SCRIPT" = true ] && [ ! $TRAIN = "na" ]; then
-    if [ -d "data/train" ] && [ "$(ls -A data/train)" ]; then
-        CONT_SCRIPT=false
-    else
-        if [ ! -d "data/train" ]; then
-            mkdir data/train
-        fi
-        if [ ! -d "data/test" ]; then
-            mkdir data/test
-        fi
-
-        cd ../utils
-
-        # Defaults to -1 if not specified, causes script to randomly generate seed
-        SEED_ARGUMENT="${SPLIT_SEED:--1}"
-
-        if [ -z $TRAIN ]; then
-            LEAF_DATA_META_DIR=${META_DIR} python3 split_data.py $NAMETAG $TFRACTAG --seed ${SEED_ARGUMENT}
-        elif [ $TRAIN = "user" ]; then
-            LEAF_DATA_META_DIR=${META_DIR} python3 split_data.py $NAMETAG --by_user $TFRACTAG --seed ${SEED_ARGUMENT}
-        elif [ $TRAIN = "sample" ]; then
-            LEAF_DATA_META_DIR=${META_DIR} python3 split_data.py $NAMETAG --by_sample $TFRACTAG --seed ${SEED_ARGUMENT}
-        fi
-
-        cd ../$NAME
-    fi
-fi
-
-if [ -z "${NO_CHECKSUM}" ]; then
-    echo '------------------------------'
-    echo "calculating JSON file checksums"
-    find 'data/' -type f -name '*.json' -exec md5sum {} + | sort -k 2 > ${CHECKSUM_FNAME}
-    echo "checksums written to ${CHECKSUM_FNAME}"
-fi
-
-if [ "$CONT_SCRIPT" = false ]; then
-    echo "Data for one of the specified preprocessing tasks has already been"
-    echo "generated. If you would like to re-generate data for this directory,"
-    echo "please delete the existing one. Otherwise, please remove the"
-    echo "respective tag(s) from the preprocessing command."
-fi
diff --git a/src/blades/models/utils/remove_users.py b/src/blades/models/utils/remove_users.py
deleted file mode 100644
index f9cf2b1..0000000
--- a/src/blades/models/utils/remove_users.py
+++ /dev/null
@@ -1,78 +0,0 @@
-'''
-removes users with less than the given number of samples
-'''
-
-import argparse
-import json
-import os
-
-from constants import DATASETS
-
-parser = argparse.ArgumentParser()
-
-parser.add_argument('--name',
-                    help='name of dataset to parse; default: sent140;',
-                    type=str,
-                    choices=DATASETS,
-                    default='FEMNIST')
-
-parser.add_argument('--min_samples',
-                    help='users with less than x samples are discarded; default: 10;',
-                    type=int,
-                    default=400)
-
-args = parser.parse_args()
-
-print('------------------------------')
-print('removing users with less than %d samples' % args.min_samples)
-
-parent_path = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
-dir = os.path.join(parent_path, args.name, 'data')
-subdir = os.path.join(dir, 'sampled_data')
-files = []
-if os.path.exists(subdir):
-    files = os.listdir(subdir)
-if len(files) == 0:
-    subdir = os.path.join(dir, 'all_data')
-    files = os.listdir(subdir)
-files = [f for f in files if f.endswith('.json')]
-
-for f in files:
-    users = []
-    hierarchies = []
-    num_samples = []
-    user_data = {}
-    
-    file_dir = os.path.join(subdir, f)
-    with open(file_dir, 'r') as inf:
-        data = json.load(inf)
-    
-    num_users = len(data['users'])
-    for i in range(num_users):
-        curr_user = data['users'][i]
-        curr_hierarchy = None
-        if 'hierarchies' in data:
-            curr_hierarchy = data['hierarchies'][i]
-        curr_num_samples = data['num_samples'][i]
-        if (curr_num_samples >= args.min_samples):
-            user_data[curr_user] = data['user_data'][curr_user]
-            users.append(curr_user)
-            if curr_hierarchy is not None:
-                hierarchies.append(curr_hierarchy)
-            num_samples.append(data['num_samples'][i])
-    
-    if not users:
-        continue
-    all_data = {}
-    all_data['users'] = users
-    if len(hierarchies) == len(users):
-        all_data['hierarchies'] = hierarchies
-    all_data['num_samples'] = num_samples
-    all_data['user_data'] = user_data
-    
-    file_name = '%s_keep_%d.json' % ((f[:-5]), args.min_samples)
-    ouf_dir = os.path.join(dir, 'rem_user_data', file_name)
-    
-    print('writing %s' % file_name)
-    with open(ouf_dir, 'w') as outfile:
-        json.dump(all_data, outfile)
diff --git a/src/blades/models/utils/sample.py b/src/blades/models/utils/sample.py
deleted file mode 100644
index 72c325b..0000000
--- a/src/blades/models/utils/sample.py
+++ /dev/null
@@ -1,200 +0,0 @@
-'''
-samples from all raw data;
-by default samples in a non-iid manner; namely, randomly selects users from 
-raw data until their cumulative amount of data exceeds the given number of 
-datapoints to sample (specified by --fraction argument);
-ordering of original data points is not preserved in sampled data
-'''
-
-import argparse
-import json
-import os
-import random
-import time
-from collections import OrderedDict
-
-from constants import DATASETS, SEED_FILES
-from util import iid_divide
-
-parser = argparse.ArgumentParser()
-
-parser.add_argument('--name',
-                    help='name of dataset to parse; default: sent140;',
-                    type=str,
-                    choices=DATASETS,
-                    default='sent140')
-parser.add_argument('--iid',
-                    help='sample iid;',
-                    action="store_true")
-parser.add_argument('--niid',
-                    help="sample niid;",
-                    dest='iid', action='store_false')
-parser.add_argument('--fraction',
-                    help='fraction of all data to sample; default: 0.1;',
-                    type=float,
-                    default=0.1)
-parser.add_argument('--u',
-                    help=('number of users in iid data set; ignored in niid case;'
-                          'represented as fraction of original total number of users; '
-                          'default: 0.01;'),
-                    type=float,
-                    default=0.01)
-parser.add_argument('--seed',
-                    help='seed for random sampling of data',
-                    type=int,
-                    default=None)
-parser.set_defaults(iid=False)
-
-args = parser.parse_args()
-
-print('------------------------------')
-print('sampling data')
-
-parent_path = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
-data_dir = os.path.join(parent_path, args.name, 'data')
-subdir = os.path.join(data_dir, 'all_data')
-files = os.listdir(subdir)
-files = [f for f in files if f.endswith('.json')]
-
-rng_seed = (args.seed if (args.seed is not None and args.seed >= 0) else int(time.time()))
-print("Using seed {}".format(rng_seed))
-rng = random.Random(rng_seed)
-print(os.environ.get('LEAF_DATA_META_DIR'))
-if os.environ.get('LEAF_DATA_META_DIR') is not None:
-    seed_fname = os.path.join(os.environ.get('LEAF_DATA_META_DIR'), SEED_FILES['sampling'])
-    with open(seed_fname, 'w+') as f:
-        f.write("# sampling_seed used by sampling script - supply as "
-                "--smplseed to preprocess.sh or --seed to utils/sample.py\n")
-        f.write(str(rng_seed))
-    print("- random seed written out to {file}".format(file=seed_fname))
-else:
-    print("- using random seed '{seed}' for sampling".format(seed=rng_seed))
-
-new_user_count = 0  # for iid case
-for f in files:
-    file_dir = os.path.join(subdir, f)
-    with open(file_dir, 'r') as inf:
-        # Load data into an OrderedDict, to prevent ordering changes
-        # and enable reproducibility
-        data = json.load(inf, object_pairs_hook=OrderedDict)
-    
-    num_users = len(data['users'])
-    
-    tot_num_samples = sum(data['num_samples'])
-    num_new_samples = int(args.fraction * tot_num_samples)
-    
-    hierarchies = None
-    
-    if (args.iid):
-        raw_list = list(data['user_data'].values())
-        raw_x = [elem['x'] for elem in raw_list]
-        raw_y = [elem['y'] for elem in raw_list]
-        x_list = [item for sublist in raw_x for item in sublist]  # flatten raw_x
-        y_list = [item for sublist in raw_y for item in sublist]  # flatten raw_y
-        
-        num_new_users = int(round(args.u * num_users))
-        if num_new_users == 0:
-            num_new_users += 1
-        
-        indices = [i for i in range(tot_num_samples)]
-        new_indices = rng.sample(indices, num_new_samples)
-        users = [str(i + new_user_count) for i in range(num_new_users)]
-        
-        user_data = {}
-        for user in users:
-            user_data[user] = {'x': [], 'y': []}
-        all_x_samples = [x_list[i] for i in new_indices]
-        all_y_samples = [y_list[i] for i in new_indices]
-        x_groups = iid_divide(all_x_samples, num_new_users)
-        y_groups = iid_divide(all_y_samples, num_new_users)
-        for i in range(num_new_users):
-            user_data[users[i]]['x'] = x_groups[i]
-            user_data[users[i]]['y'] = y_groups[i]
-        
-        num_samples = [len(user_data[u]['y']) for u in users]
-        
-        new_user_count += num_new_users
-    
-    else:
-        
-        ctot_num_samples = 0
-        
-        users = data['users']
-        users_and_hiers = None
-        if 'hierarchies' in data:
-            users_and_hiers = list(zip(users, data['hierarchies']))
-            rng.shuffle(users_and_hiers)
-        else:
-            rng.shuffle(users)
-        user_i = 0
-        num_samples = []
-        user_data = {}
-        
-        if 'hierarchies' in data:
-            hierarchies = []
-        
-        while (ctot_num_samples < num_new_samples):
-            hierarchy = None
-            if users_and_hiers is not None:
-                user, hier = users_and_hiers[user_i]
-            else:
-                user = users[user_i]
-            
-            cdata = data['user_data'][user]
-            
-            cnum_samples = len(data['user_data'][user]['y'])
-            
-            if (ctot_num_samples + cnum_samples > num_new_samples):
-                cnum_samples = num_new_samples - ctot_num_samples
-                indices = [i for i in range(cnum_samples)]
-                new_indices = rng.sample(indices, cnum_samples)
-                x = []
-                y = []
-                for i in new_indices:
-                    x.append(data['user_data'][user]['x'][i])
-                    y.append(data['user_data'][user]['y'][i])
-                cdata = {'x': x, 'y': y}
-            
-            if 'hierarchies' in data:
-                hierarchies.append(hier)
-            
-            num_samples.append(cnum_samples)
-            user_data[user] = cdata
-            
-            ctot_num_samples += cnum_samples
-            user_i += 1
-        
-        if 'hierarchies' in data:
-            users = [u for u, h in users_and_hiers][:user_i]
-        else:
-            users = users[:user_i]
-    
-    # ------------
-    # create .json file
-    
-    all_data = {}
-    all_data['users'] = users
-    if hierarchies is not None:
-        all_data['hierarchies'] = hierarchies
-    all_data['num_samples'] = num_samples
-    all_data['user_data'] = user_data
-    
-    slabel = ''
-    if (args.iid):
-        slabel = 'iid'
-    else:
-        slabel = 'niid'
-    
-    arg_frac = str(args.fraction)
-    arg_frac = arg_frac[2:]
-    arg_nu = str(args.u)
-    arg_nu = arg_nu[2:]
-    arg_label = arg_frac
-    if (args.iid):
-        arg_label = '%s_%s' % (arg_nu, arg_label)
-    file_name = '%s_%s_%s.json' % ((f[:-5]), slabel, arg_label)
-    ouf_dir = os.path.join(data_dir, 'sampled_data', file_name)
-    
-    print('writing %s' % file_name)
-    with open(ouf_dir, 'w') as outfile:
-        json.dump(all_data, outfile)
diff --git a/src/blades/models/utils/split_data.py b/src/blades/models/utils/split_data.py
deleted file mode 100644
index a6acd6f..0000000
--- a/src/blades/models/utils/split_data.py
+++ /dev/null
@@ -1,269 +0,0 @@
-'''
-splits data into train and test sets
-'''
-
-import argparse
-import json
-import os
-import random
-import sys
-import time
-from collections import OrderedDict
-
-from constants import DATASETS, SEED_FILES
-
-
-def create_jsons_for(user_files, which_set, max_users, include_hierarchy):
-    '''used in split-by-user case'''
-    user_count = 0
-    json_index = 0
-    users = []
-    if include_hierarchy:
-        hierarchies = []
-    else:
-        hierarchies = None
-    num_samples = []
-    user_data = {}
-    for (i, t) in enumerate(user_files):
-        if include_hierarchy:
-            (u, h, ns, f) = t
-        else:
-            (u, ns, f) = t
-        
-        file_dir = os.path.join(subdir, f)
-        with open(file_dir, 'r') as inf:
-            data = json.load(inf)
-        
-        users.append(u)
-        if include_hierarchy:
-            hierarchies.append(h)
-        num_samples.append(ns)
-        user_data[u] = data['user_data'][u]
-        user_count += 1
-        
-        if (user_count == max_users) or (i == len(user_files) - 1):
-            
-            all_data = {}
-            all_data['users'] = users
-            if include_hierarchy:
-                all_data['hierarchies'] = hierarchies
-            all_data['num_samples'] = num_samples
-            all_data['user_data'] = user_data
-            
-            data_i = f.find('data')
-            num_i = data_i + 5
-            num_to_end = f[num_i:]
-            param_i = num_to_end.find('_')
-            param_to_end = '.json'
-            if param_i != -1:
-                param_to_end = num_to_end[param_i:]
-            nf = '%s_%d%s' % (f[:(num_i - 1)], json_index, param_to_end)
-            file_name = '%s_%s_%s.json' % ((nf[:-5]), which_set, arg_label)
-            ouf_dir = os.path.join(dir, which_set, file_name)
-            
-            print('writing %s' % file_name)
-            with open(ouf_dir, 'w') as outfile:
-                json.dump(all_data, outfile)
-            
-            user_count = 0
-            json_index += 1
-            users = []
-            if include_hierarchy:
-                hierarchies = []
-            num_samples = []
-            user_data = {}
-
-
-parser = argparse.ArgumentParser()
-
-parser.add_argument('--name',
-                    help='name of dataset to parse; default: sent140;',
-                    type=str,
-                    choices=DATASETS,
-                    default='femnist')
-parser.add_argument('--by_user',
-                    help='divide users into training and test set groups;',
-                    dest='user', action='store_true')
-parser.add_argument('--by_sample',
-                    help="divide each user's samples into training and test set groups;",
-                    dest='user', action='store_false')
-parser.add_argument('--frac',
-                    help='fraction in training set; default: 0.9;',
-                    type=float,
-                    default=0.8)
-parser.add_argument('--seed',
-                    help='seed for random partitioning of test/train data',
-                    type=int,
-                    default=None)
-
-parser.set_defaults(user=False)
-
-args = parser.parse_args()
-
-print('------------------------------')
-print('generating training and test sets')
-
-parent_path = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
-dir = os.path.join(parent_path, args.name, 'data')
-subdir = os.path.join(dir, 'rem_user_data')
-files = []
-if os.path.exists(subdir):
-    files = os.listdir(subdir)
-if len(files) == 0:
-    subdir = os.path.join(dir, 'sampled_data')
-    if os.path.exists(subdir):
-        files = os.listdir(subdir)
-if len(files) == 0:
-    subdir = os.path.join(dir, 'all_data')
-    files = os.listdir(subdir)
-files = [f for f in files if f.endswith('.json')]
-
-rng_seed = (args.seed if (args.seed is not None and args.seed >= 0) else int(time.time()))
-rng = random.Random(rng_seed)
-if os.environ.get('LEAF_DATA_META_DIR') is not None:
-    seed_fname = os.path.join(os.environ.get('LEAF_DATA_META_DIR'), SEED_FILES['split'])
-    with open(seed_fname, 'w+') as f:
-        f.write("# split_seed used by sampling script - supply as "
-                "--spltseed to preprocess.sh or --seed to utils/split_data.py\n")
-        f.write(str(rng_seed))
-    print("- random seed written out to {file}".format(file=seed_fname))
-else:
-    print("- using random seed '{seed}' for sampling".format(seed=rng_seed))
-
-arg_label = str(args.frac)
-arg_label = arg_label[2:]
-
-# check if data contains information on hierarchies
-file_dir = os.path.join(subdir, files[0])
-with open(file_dir, 'r') as inf:
-    data = json.load(inf)
-include_hierarchy = 'hierarchies' in data
-
-if (args.user):
-    print('splitting data by user')
-    
-    # 1 pass through all the json files to instantiate arr
-    # containing all possible (user, .json file name) tuples
-    user_files = []
-    for f in files:
-        file_dir = os.path.join(subdir, f)
-        with open(file_dir, 'r') as inf:
-            # Load data into an OrderedDict, to prevent ordering changes
-            # and enable reproducibility
-            data = json.load(inf, object_pairs_hook=OrderedDict)
-        if include_hierarchy:
-            user_files.extend([(u, h, ns, f) for (u, h, ns) in
-                               zip(data['users'], data['hierarchies'], data['num_samples'])])
-        else:
-            user_files.extend([(u, ns, f) for (u, ns) in
-                               zip(data['users'], data['num_samples'])])
-    
-    # randomly sample from user_files to pick training set users
-    num_users = len(user_files)
-    num_train_users = int(args.frac * num_users)
-    indices = [i for i in range(num_users)]
-    train_indices = rng.sample(indices, num_train_users)
-    train_blist = [False for i in range(num_users)]
-    for i in train_indices:
-        train_blist[i] = True
-    train_user_files = []
-    test_user_files = []
-    for i in range(num_users):
-        if (train_blist[i]):
-            train_user_files.append(user_files[i])
-        else:
-            test_user_files.append(user_files[i])
-    
-    max_users = sys.maxsize
-    if args.name == 'femnist':
-        max_users = 50  # max number of users per json file
-    create_jsons_for(train_user_files, 'train', max_users, include_hierarchy)
-    create_jsons_for(test_user_files, 'test', max_users, include_hierarchy)
-
-else:
-    print('splitting data by sample')
-    
-    for f in files:
-        file_dir = os.path.join(subdir, f)
-        with open(file_dir, 'r') as inf:
-            # Load data into an OrderedDict, to prevent ordering changes
-            # and enable reproducibility
-            data = json.load(inf, object_pairs_hook=OrderedDict)
-        
-        num_samples_train = []
-        user_data_train = {}
-        num_samples_test = []
-        user_data_test = {}
-        
-        user_indices = []  # indices of users in data['users'] that are not deleted
-        
-        removed = 0
-        for i, u in enumerate(data['users']):
-            
-            curr_num_samples = len(data['user_data'][u]['y'])
-            if curr_num_samples >= 2:
-                # ensures number of train and test samples both >= 1
-                num_train_samples = max(1, int(args.frac * curr_num_samples))
-                if curr_num_samples == 2:
-                    num_train_samples = 1
-                
-                num_test_samples = curr_num_samples - num_train_samples
-                
-                indices = [j for j in range(curr_num_samples)]
-                if args.name in ['shakespeare']:
-                    train_indices = [i for i in range(num_train_samples)]
-                    test_indices = [i for i in range(num_train_samples + 80 - 1, curr_num_samples)]
-                else:
-                    train_indices = rng.sample(indices, num_train_samples)
-                    test_indices = [i for i in range(curr_num_samples) if i not in train_indices]
-                
-                if len(train_indices) >= 1 and len(test_indices) >= 1:
-                    user_indices.append(i)
-                    num_samples_train.append(num_train_samples)
-                    num_samples_test.append(num_test_samples)
-                    user_data_train[u] = {'x': [], 'y': []}
-                    user_data_test[u] = {'x': [], 'y': []}
-                    
-                    train_blist = [False for _ in range(curr_num_samples)]
-                    test_blist = [False for _ in range(curr_num_samples)]
-                    
-                    for j in train_indices:
-                        train_blist[j] = True
-                    for j in test_indices:
-                        test_blist[j] = True
-                    
-                    for j in range(curr_num_samples):
-                        if (train_blist[j]):
-                            user_data_train[u]['x'].append(data['user_data'][u]['x'][j])
-                            user_data_train[u]['y'].append(data['user_data'][u]['y'][j])
-                        elif (test_blist[j]):
-                            user_data_test[u]['x'].append(data['user_data'][u]['x'][j])
-                            user_data_test[u]['y'].append(data['user_data'][u]['y'][j])
-        
-        users = [data['users'][i] for i in user_indices]
-        
-        all_data_train = {}
-        all_data_train['users'] = users
-        all_data_train['num_samples'] = num_samples_train
-        all_data_train['user_data'] = user_data_train
-        
-        all_data_test = {}
-        all_data_test['users'] = users
-        all_data_test['num_samples'] = num_samples_test
-        all_data_test['user_data'] = user_data_test
-        
-        if include_hierarchy:
-            hierarchies = [data['hierarchies'][i] for i in user_indices]
-            all_data_train['hierarchies'] = hierarchies
-            all_data_test['hierarchies'] = hierarchies
-        
-        file_name_train = '%s_train_%s.json' % ((f[:-5]), arg_label)
-        file_name_test = '%s_test_%s.json' % ((f[:-5]), arg_label)
-        ouf_dir_train = os.path.join(dir, 'train', file_name_train)
-        ouf_dir_test = os.path.join(dir, 'test', file_name_test)
-        print('writing %s' % file_name_train)
-        with open(ouf_dir_train, 'w') as outfile:
-            json.dump(all_data_train, outfile)
-        print('writing %s' % file_name_test)
-        with open(ouf_dir_test, 'w') as outfile:
-            json.dump(all_data_test, outfile)
diff --git a/src/blades/models/utils/stats.py b/src/blades/models/utils/stats.py
deleted file mode 100644
index 0428c50..0000000
--- a/src/blades/models/utils/stats.py
+++ /dev/null
@@ -1,92 +0,0 @@
-'''
-assumes that the user has already generated .json file(s) containing data
-'''
-
-import argparse
-import json
-import os
-
-import matplotlib.pyplot as plt
-import numpy as np
-from scipy import stats
-
-from constants import DATASETS
-
-parser = argparse.ArgumentParser()
-
-parser.add_argument('--name',
-                    help='name of dataset to parse; default: sent140;',
-                    type=str,
-                    choices=DATASETS,
-                    default='sent140')
-
-args = parser.parse_args()
-
-
-def load_data(name):
-    users = []
-    num_samples = []
-    
-    parent_path = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
-    data_dir = os.path.join(parent_path, name, 'data')
-    subdir = os.path.join(data_dir, 'all_data')
-    
-    files = os.listdir(subdir)
-    files = [f for f in files if f.endswith('.json')]
-    
-    for f in files:
-        file_dir = os.path.join(subdir, f)
-        
-        with open(file_dir) as inf:
-            data = json.load(inf)
-        
-        users.extend(data['users'])
-        num_samples.extend(data['num_samples'])
-    
-    return users, num_samples
-
-
-def print_dataset_stats(name):
-    users, num_samples = load_data(name)
-    num_users = len(users)
-    
-    print('####################################')
-    print('DATASET: %s' % name)
-    print('%d users' % num_users)
-    print('%d samples (total)' % np.sum(num_samples))
-    print('%.2f samples per user (mean)' % np.mean(num_samples))
-    print('num_samples (std): %.2f' % np.std(num_samples))
-    print('num_samples (std/mean): %.2f' % (np.std(num_samples) / np.mean(num_samples)))
-    print('num_samples (skewness): %.2f' % stats.skew(num_samples))
-    
-    bins = [0, 20, 40, 60, 80, 100, 120, 140, 160, 180, 200]
-    if args.name == 'shakespeare':
-        bins = [0, 2000, 4000, 6000, 8000, 10000, 12000, 14000, 16000, 18000, 20000]
-    if args.name == 'femnist':
-        bins = [0, 20, 40, 60, 80, 100, 120, 140, 160, 180, 200, 220, 240, 260, 280, 300, 320, 340, 360, 380, 400, 420,
-                440, 460, 480, 500]
-    if args.name == 'celeba':
-        bins = [2 * i for i in range(20)]
-    if args.name == 'sent140':
-        bins = [i for i in range(16)]
-    
-    hist, edges = np.histogram(num_samples, bins=bins)
-    for e, h in zip(edges, hist):
-        print(e, "\t", h)
-    
-    parent_path = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
-    data_dir = os.path.join(parent_path, name, 'data')
-    
-    plt.hist(num_samples, bins=bins)
-    fig_name = "%s_hist_nolabel.png" % name
-    fig_dir = os.path.join(data_dir, fig_name)
-    plt.savefig(fig_dir)
-    plt.title(name)
-    plt.xlabel('number of samples')
-    plt.ylabel("number of users")
-    fig_name = "%s_hist.png" % name
-    fig_dir = os.path.join(data_dir, fig_name)
-    plt.savefig(fig_dir)
-
-
-print_dataset_stats(args.name)
diff --git a/src/blades/models/utils/util.py b/src/blades/models/utils/util.py
index f8286d0..9268a41 100644
--- a/src/blades/models/utils/util.py
+++ b/src/blades/models/utils/util.py
@@ -2,30 +2,10 @@
 
 
 def save_obj(obj, name):
-    with open(name + '.pkl', 'wb') as f:
+    with open(name + ".pkl", "wb") as f:
         pickle.dump(obj, f, pickle.HIGHEST_PROTOCOL)
 
 
 def load_obj(name):
-    with open(name + '.pkl', 'rb') as f:
+    with open(name + ".pkl", "rb") as f:
         return pickle.load(f)
-
-
-def iid_divide(l, g):
-    '''
-    divide list l among g groups
-    each group has either int(len(l)/g) or int(len(l)/g)+1 elements
-    returns a list of groups
-    '''
-    num_elems = len(l)
-    group_size = int(len(l) / g)
-    num_big_groups = num_elems - g * group_size
-    num_small_groups = g - num_big_groups
-    glist = []
-    for i in range(num_small_groups):
-        glist.append(l[group_size * i: group_size * (i + 1)])
-    bi = group_size * num_small_groups
-    group_size += 1
-    for i in range(num_big_groups):
-        glist.append(l[bi + group_size * i:bi + group_size * (i + 1)])
-    return glist
diff --git a/src/blades/server.py b/src/blades/server.py
deleted file mode 100644
index b4ddd56..0000000
--- a/src/blades/server.py
+++ /dev/null
@@ -1,75 +0,0 @@
-from typing import Callable
-
-import torch
-
-
-class BladesServer(object):
-    r'''Simulating the server of the federated learning system.
-        
-        :ivar aggregator: a callable which takes a list of tensors and returns
-                an aggregated tensor.
-        :vartype aggregator: callable
-    
-        :param  optimizer: The global optimizer, which can be any optimizer from Pytorch.
-        :type optimizer: torch.optim.Optimizer
-        :param model: The global model
-        :type model: torch.nn.Module
-        :param aggregator: a callable which takes a list of tensors and returns
-                an aggregated tensor.
-        :type aggregator: callable
-    '''
-    
-    def __init__(self,
-                 optimizer: torch.optim.Optimizer,
-                 model: torch.nn.Module,
-                 aggregator: Callable[[list], torch.Tensor],
-                 *args,
-                 **kwargs,
-                 ):
-        self.optimizer = optimizer
-        self.model = model
-        self.aggregator = aggregator
-    
-    def get_opt(self) -> torch.optim.Optimizer:
-        r'''
-        Returns the global optimizer.
-        '''
-        return self.optimizer
-    
-    def zero_grad(self, set_to_none: bool = False):
-        r"""Sets the gradients of all optimized :class:`torch.Tensor` s to zero.
-        It should be called before assigning pseudo-gradient
-
-        :param set_to_none: See `Pytorch documentation <https://pytorch.org/docs/stable/generated/torch.optim.Optimizer.zero_grad.html>`_.
-        :type set_to_none: bool
-        """
-        self.optimizer.zero_grad(set_to_none=set_to_none)
-    
-    def get_model(self) -> torch.nn.Module:
-        r'''
-        Returns the current global model.
-        '''
-        return self.model
-    
-    def apply_update(self, update: torch.Tensor) -> None:
-        r'''
-        Apply a step of global optimization.
-            
-            .. note::
-                The input should be a ``Tensor``, which will be converted to ``pseudo-gradient``
-                layer by layer.
-                
-        :param update: The aggregated update.
-        :type update: torch.Tensor
-        '''
-        self.zero_grad()
-        beg = 0
-        for group in self.optimizer.param_groups:
-            for p in group["params"]:
-                if not p.requires_grad:
-                    continue
-                end = beg + len(p.data.view(-1))
-                x = update[beg:end].reshape_as(p.data)
-                p.grad = -x.clone().detach().to(p.device)
-                beg = end
-        self.optimizer.step()
diff --git a/src/blades/utils/torch_utils.py b/src/blades/utils/torch_utils.py
index 0d1db1f..b1571d3 100644
--- a/src/blades/utils/torch_utils.py
+++ b/src/blades/utils/torch_utils.py
@@ -1,4 +1,4 @@
-from typing import Union, Iterable
+from typing import Iterable, Union
 
 import torch
 import torch.nn as nn
@@ -12,7 +12,7 @@
 class HLoss(nn.Module):
     def __init__(self):
         super(HLoss, self).__init__()
-    
+
     def forward(self, x):
         b = F.softmax(x, dim=1) * F.log_softmax(x, dim=1)
         b = -1.0 * b.sum()
@@ -20,45 +20,70 @@ def forward(self, x):
 
 
 def l2dist(model1, model2):
-    return LA.norm(torch.tensor([LA.norm(model1[k] - model2[k]) for k in model1 if model1[k].dtype != torch.int64]))
+    return LA.norm(
+        torch.tensor(
+            [
+                LA.norm(model1[k] - model2[k])
+                for k in model1
+                if model1[k].dtype != torch.int64
+            ]
+        )
+    )
 
 
 def l2norm(model):
-    return LA.norm(torch.tensor([LA.norm(model[k]) for k in model if model[k].dtype != torch.int64]))
+    return LA.norm(
+        torch.tensor(
+            [LA.norm(model[k]) for k in model if model[k].dtype != torch.int64]
+        )
+    )
 
 
 def cos_sim(model1, model2):
-    return torch.sum(torch.tensor([torch.sum(model1[k] * model2[k]) for k in model1])) / torch.max(
-        l2norm(model1) * l2norm(model2), torch.tensor(0.00001))
+    return torch.sum(
+        torch.tensor([torch.sum(model1[k] * model2[k]) for k in model1])
+    ) / torch.max(l2norm(model1) * l2norm(model2), torch.tensor(0.00001))
 
 
 def clip_para_norm_(
-        parameters: _tensor_or_tensors, max_norm: float, norm_type: float = 2.0,
-        error_if_nonfinite: bool = False) -> torch.Tensor:
+    parameters: _tensor_or_tensors,
+    max_norm: float,
+    norm_type: float = 2.0,
+    error_if_nonfinite: bool = False,
+) -> torch.Tensor:
     if isinstance(parameters, torch.Tensor):
         parameters = [parameters]
     parameters = [p for p in parameters.values()]
     max_norm = float(max_norm)
     norm_type = float(norm_type)
     if len(parameters) == 0:
-        return torch.tensor(0.)
+        return torch.tensor(0.0)
     device = parameters[0].device
     if norm_type == inf:
         norms = [p.detach().abs().max().to(device) for p in parameters]
         total_norm = norms[0] if len(norms) == 1 else torch.max(torch.stack(norms))
     else:
-        total_norm = torch.norm(torch.stack([torch.norm(p.detach(), norm_type).to(device) for p in parameters
-                                             if p.dtype != torch.int64]), norm_type)
+        total_norm = torch.norm(
+            torch.stack(
+                [
+                    torch.norm(p.detach(), norm_type).to(device)
+                    for p in parameters
+                    if p.dtype != torch.int64
+                ]
+            ),
+            norm_type,
+        )
     if error_if_nonfinite and torch.logical_or(total_norm.isnan(), total_norm.isinf()):
         raise RuntimeError(
-            f'The total norm of order {norm_type} for gradients from '
-            '`parameters` is non-finite, so it cannot be clipped. To disable '
-            'this error and scale the gradients by the non-finite norm anyway, '
-            'set `error_if_nonfinite=False`')
+            f"The total norm of order {norm_type} for gradients from "
+            "`parameters` is non-finite, so it cannot be clipped. To disable "
+            "this error and scale the gradients by the non-finite norm anyway, "
+            "set `error_if_nonfinite=False`"
+        )
     clip_coef = max_norm / (total_norm + 1e-6)
-    # Note: multiplying by the clamped coef is redundant when the coef is clamped to 1, but doing so
-    # avoids a `if clip_coef < 1:` conditional which can require a CPU <=> device synchronization
-    # when the gradients do not reside in CPU memory.
+    # Note: multiplying by the clamped coef is redundant when the coef is clamped to 1,
+    # but doing so avoids a `if clip_coef < 1:` conditional which can require a CPU <=>
+    # device synchronization when the gradients do not reside in CPU memory.
     clip_coef_clamped = torch.clamp(clip_coef, max=1.0)
     for p in parameters:
         if p.dtype != torch.int64:
@@ -67,31 +92,43 @@ def clip_para_norm_(
 
 
 def clip_tensor_norm_(
-        parameters: _tensor_or_tensors, max_norm: float, norm_type: float = 2.0,
-        error_if_nonfinite: bool = False) -> torch.Tensor:
+    parameters: _tensor_or_tensors,
+    max_norm: float,
+    norm_type: float = 2.0,
+    error_if_nonfinite: bool = False,
+) -> torch.Tensor:
     if isinstance(parameters, torch.Tensor):
         parameters = [parameters]
     max_norm = float(max_norm)
     norm_type = float(norm_type)
     if len(parameters) == 0:
-        return torch.tensor(0.)
+        return torch.tensor(0.0)
     device = parameters[0].device
     if norm_type == inf:
         norms = [p.detach().abs().max().to(device) for p in parameters]
         total_norm = norms[0] if len(norms) == 1 else torch.max(torch.stack(norms))
     else:
-        total_norm = torch.norm(torch.stack([torch.norm(p.detach(), norm_type).to(device) for p in parameters
-                                             if p.dtype != torch.int64]), norm_type)
+        total_norm = torch.norm(
+            torch.stack(
+                [
+                    torch.norm(p.detach(), norm_type).to(device)
+                    for p in parameters
+                    if p.dtype != torch.int64
+                ]
+            ),
+            norm_type,
+        )
     if error_if_nonfinite and torch.logical_or(total_norm.isnan(), total_norm.isinf()):
         raise RuntimeError(
-            f'The total norm of order {norm_type} for gradients from '
-            '`parameters` is non-finite, so it cannot be clipped. To disable '
-            'this error and scale the gradients by the non-finite norm anyway, '
-            'set `error_if_nonfinite=False`')
+            f"The total norm of order {norm_type} for gradients from "
+            "`parameters` is non-finite, so it cannot be clipped. To disable "
+            "this error and scale the gradients by the non-finite norm anyway, "
+            "set `error_if_nonfinite=False`"
+        )
     clip_coef = max_norm / (total_norm + 1e-6)
-    # Note: multiplying by the clamped coef is redundant when the coef is clamped to 1, but doing so
-    # avoids a `if clip_coef < 1:` conditional which can require a CPU <=> device synchronization
-    # when the gradients do not reside in CPU memory.
+    # Note: multiplying by the clamped coef is redundant when the coef is clamped to 1,
+    # but doing so avoids a `if clip_coef < 1:` conditional which can require a CPU <=>
+    # device synchronization when the gradients do not reside in CPU memory.
     clip_coef_clamped = torch.clamp(clip_coef, max=1.0)
     for p in parameters:
         if p.dtype != torch.int64:
diff --git a/src/blades/utils/utils.py b/src/blades/utils/utils.py
index dfd002e..498a2ee 100644
--- a/src/blades/utils/utils.py
+++ b/src/blades/utils/utils.py
@@ -5,9 +5,10 @@
 from importlib import reload
 
 import numpy as np
+import ruamel.yaml as yaml
 import torch
 from torch import nn
-import ruamel.yaml as yaml
+
 
 class BColors(object):
     HEADER = "\033[95m"
@@ -21,7 +22,6 @@ class BColors(object):
     UNDERLINE = "\033[4m"
 
 
-
 def touch(fname: str, times=None, create_dirs: bool = False):
     if create_dirs:
         base_dir = os.path.dirname(fname)
@@ -37,14 +37,14 @@ def touch_dir(base_dir: str) -> None:
 
 
 def accuracy(output, target, topk=(1,)):
-    """Computes the precision@k for the specified values of k"""
+    """Computes the precision@k for the specified values of k."""
     maxk = max(topk)
     batch_size = target.size(0)
-    
+
     _, pred = output.topk(maxk, 1, True, True)
     pred = pred.t()
     correct = pred.eq(target.view(1, -1).expand_as(pred))
-    
+
     res = []
     for k in topk:
         correct_k = correct[:k].view(-1).float().sum(0)
@@ -72,9 +72,9 @@ def initialize_logger(log_root):
     else:
         shutil.rmtree(log_root)
         os.makedirs(log_root)
-    
+
     print(f"Logging files to {log_root}")
-    
+
     # Only to file; One dict per line; Easy to process
     json_logger = logging.getLogger("stats")
     json_logger.setLevel(logging.INFO)
@@ -82,7 +82,7 @@ def initialize_logger(log_root):
     fh.setLevel(logging.INFO)
     fh.setFormatter(logging.Formatter("%(message)s"))
     json_logger.addHandler(fh)
-    
+
     debug_logger = logging.getLogger("debug")
     debug_logger.setLevel(logging.INFO)
     # ch = logging.StreamHandler()
@@ -98,39 +98,40 @@ def initialize_logger(log_root):
 def reset_model_weights(model: nn.Module) -> None:
     """
     refs:
-        - https://discuss.pytorch.org/t/how-to-re-set-alll-parameters-in-a-network/20819/6
-        - https://stackoverflow.com/questions/63627997/reset-parameters-of-a-neural-network-in-pytorch
-        - https://pytorch.org/docs/stable/generated/torch.nn.Module.html
+
+    - https://discuss.pytorch.org/t/how-to-re-set-alll-parameters-in-a-network/20819/6
+    - https://pytorch.org/docs/stable/generated/torch.nn.Module.html
     """
 
     @torch.no_grad()
     def weight_reset(m: nn.Module):
-        # - check if the current module has reset_parameters & if it's callabed called it on m
+        # - check if the current module has reset_parameters & if it's callabed
+        # called it on m
         reset_parameters = getattr(m, "reset_parameters", None)
         if callable(reset_parameters):
             m.reset_parameters()
 
-    # Applies fn recursively to every submodule see: https://pytorch.org/docs/stable/generated/torch.nn.Module.html
+    # Applies fn recursively to every submodule see:
+    # https://pytorch.org/docs/stable/generated/torch.nn.Module.html
     model.apply(fn=weight_reset)
 
+
 def set_random_seed(seed_value=0, use_cuda=False):
     np.random.seed(seed_value)  # cpu vars
     random.seed(seed_value)  # Python
     torch.manual_seed(seed_value)
-    os.environ['PYTHONHASHSEED'] = str(seed_value)  # Python hash buildin
+    os.environ["PYTHONHASHSEED"] = str(seed_value)  # Python hash buildin
     if use_cuda:
         torch.cuda.manual_seed_all(seed_value)
         torch.backends.cudnn.deterministic = True
         torch.backends.cudnn.benchmark = False
-        
-        
+
+
 def over_write_args_from_file(args, yml):
-    """
-    overwrite arguments according to config file
-    """
-    if yml == '':
+    """overwrite arguments according to config file."""
+    if yml == "":
         return
-    with open(yml, 'r', encoding='utf-8') as f:
+    with open(yml, "r", encoding="utf-8") as f:
         dic = yaml.load(f.read(), Loader=yaml.Loader)
         for k in dic:
-            setattr(args, k, dic[k])
\ No newline at end of file
+            setattr(args, k, dic[k])
diff --git a/src/setup.py b/src/setup.py
index 28eead0..c95b038 100644
--- a/src/setup.py
+++ b/src/setup.py
@@ -1,40 +1,42 @@
-from setuptools import setup, find_packages
+from setuptools import find_packages, setup
 
 requirements = [
-    'torch>=1.10.2',
-    'torchvision>=0.11.3',
-    'numpy>=1.19.4',
-    'scipy>=1.5.4',
-    'ray>=1.0.0',
-    'sklearn>=0.0',
-    'scikit-learn>=1.0.2',
-    'matplotlib>=3.4.1',
-    'requests>=2.27.1',
-    'setuptools~=58.0.4',
-    'ruamel.yaml',
+    "torch>=1.10.2",
+    "torchvision>=0.11.3",
+    "numpy>=1.19.4",
+    "scipy>=1.5.4",
+    "ray>=1.0.0",
+    "sklearn>=0.0",
+    "scikit-learn>=1.0.2",
+    "matplotlib>=3.4.1",
+    "requests>=2.27.1",
+    "setuptools~=58.0.4",
+    "ruamel.yaml",
 ]
 setup(
-    name='blades',
-    version='0.1.1',
-    description="""Blades: A simulator for Byzantine-robust federated Learning with Attacks and Defenses Experimental Simulation""",
+    name="blades",
+    version="0.1.1",
+    description="""
+    Blades: A simulator for Byzantine-robust federated Learning with Attacks and Defenses
+    """,
     # long_description=open('../README.rst').read(),
     # long_description_content_type="text/markdown",
-    author='Blades Team',
-    author_email='shenghui.li@it.uu.se',
-    url='https://bladesteam.github.io/',
+    author="Blades Team",
+    author_email="shenghui.li@it.uu.se",
+    url="https://bladesteam.github.io/",
     # py_modules=['blades'],
-    python_requires='>=3.8',
-    license='MIT license',
+    python_requires=">=3.8",
+    license="Apache License 2.0",
     zip_safe=False,
     # entry_points={
     #     'console_scripts': [""]
     # },
     install_requires=requirements,
-    keywords='Federated learning',
+    keywords="Federated learning",
     packages=find_packages(),
     classifiers=[
         "Programming Language :: Python :: 3",
-        "License :: OSI Approved :: MIT License",
+        "License :: OSI Approved :: Apache Software License",
         "Operating System :: OS Independent",
     ],
 )


Test patch:
diff --git a/src/tests/test_aggregators.py b/src/tests/test_aggregators.py
new file mode 100644
index 0000000..b2c20d5
--- /dev/null
+++ b/src/tests/test_aggregators.py
@@ -0,0 +1,38 @@
+# import pytest
+import torch
+
+from blades.aggregators import Mean, Median, Trimmedmean
+
+# from blades.aggregators import (Autogm, Clippedclustering, Clustering, Dnc,
+#                                 Fltrust, Geomed, Mean, Median, Multikrum,
+#                                 Signguard, Trimmedmean)
+
+test_data = torch.Tensor(
+    [
+        [1, 2, 3],
+        [-1, 4, -1],
+        [2.0, 2, 3],
+        [3.0, 1.0, 3],
+    ]
+)
+
+
+def test_median():
+    median = Median()
+    assert torch.equal(
+        median(test_data),
+        torch.Tensor([1.5, 2.0, 3.0000]),
+    )
+
+
+def test_mean():
+    mean = Mean()
+    assert torch.equal(
+        mean(test_data),
+        torch.Tensor([1.2500, 2.2500, 2.0000]),
+    )
+
+
+def test_trimmedmean():
+    tm = Trimmedmean(1)
+    assert torch.equal(tm(test_data), torch.Tensor([1.5000, 2.0000, 3.0000]))
diff --git a/src/tests/test_core/test_client.py b/src/tests/test_core/test_client.py
new file mode 100644
index 0000000..b74d6de
--- /dev/null
+++ b/src/tests/test_core/test_client.py
@@ -0,0 +1,7 @@
+from blades.core.client import BladesClient
+
+
+def test_base_client():
+    id = "1"
+    client = BladesClient(id=id)
+    assert client.id() == id