Repository: datamol-io/graphium (language: python)

Issue:
Fix task head test
While working on the dispatcher I noticed that some tests were broken due to the change to task heads format in the YAML files. 

This diff fixes the YAML and the tests.

Note that this PR is branching off of the dispatcher branch and should be landed after that PR lands.


Discussion / hints:

Fix patch:
diff --git a/.gitignore b/.gitignore
index 0b25e159f..7b8c7590c 100644
--- a/.gitignore
+++ b/.gitignore
@@ -16,6 +16,7 @@ lightning_logs/
 logs/
 models_checkpoints/
 outputs/
+out/
 saved_models/
 wandb/
 out/
@@ -27,6 +28,7 @@ goli/data/BindingDB/
 goli/data/mega-pubchem/
 goli/data/cache/
 goli/data/b3lyp/
+goli/data/PCQM4Mv2/
 predictions/
 
 # Others
diff --git a/README.md b/README.md
index 93f51b644..2ae0d50d3 100644
--- a/README.md
+++ b/README.md
@@ -57,16 +57,14 @@ mkdir ~/.venv                           # Create the folder for the environment
 python3 -m venv ~/.venv/goli_ipu        # Create the environment
 source ~/.venv/goli_ipu/bin/activate    # Activate the environment
 
-# Installing the dependencies for the IPU environment
-pip install torch==1.10+cpu torchvision==0.11+cpu torchaudio==0.10 -f https://download.pytorch.org/whl/torch_stable.html
-pip install torch-scatter torch-sparse torch-cluster torch-spline-conv torch-geometric -f https://data.pyg.org/whl/torch-1.10.0+cpu.html
-pip install dgl dglgo -f https://data.dgl.ai/wheels/repo.html
+# Activate poplar SDK.
+source PATH_TO_SDK/enable
 
 # Installing the poptorch SDK. Make sure to change the path
 pip install PATH_TO_SDK/poptorch-3.0.0+86945_163b7ce462_ubuntu_20_04-cp38-cp38-linux_x86_64.wh
 
-# Install the remaining requirements
-pip install -r requirements.txt
+# Install the IPU specific and goli requirements
+pip install -r requirements_ipu.txt
 
 # Install Goli in dev mode
 pip install -e .
diff --git a/docs/api/goli.trainer.md b/docs/api/goli.trainer.md
index 499888696..a0e8ed863 100644
--- a/docs/api/goli.trainer.md
+++ b/docs/api/goli.trainer.md
@@ -1,4 +1,3 @@
-::: goli.trainer.loggers
 ::: goli.trainer.metrics
 ::: goli.trainer.predictor_options
 ::: goli.trainer.predictor_summaries
diff --git a/docs/tutorials/model_training/config_ipu_tutorials.yaml b/docs/tutorials/model_training/config_ipu_tutorials.yaml
index cc90dc4d5..da4876a7b 100644
--- a/docs/tutorials/model_training/config_ipu_tutorials.yaml
+++ b/docs/tutorials/model_training/config_ipu_tutorials.yaml
@@ -238,12 +238,15 @@ predictor:
 metrics:
   homo:
     - name: mae
-      metric: mae
+      metric: mae_ipu
+      target_nan_mask: null
+      multitask_handling: flatten
       threshold_kwargs: null
     - name: pearsonr
-      metric: pearsonr
+      metric: pearsonr_ipu
       threshold_kwargs: null
-      target_nan_mask: ignore-mean-label
+      target_nan_mask: null
+      multitask_handling: mean-per-label
     # - name: f1 > 0.5
     #   metric: f1
     #   num_classes: 2
@@ -292,6 +295,6 @@ trainer:
     save_top_k: 1
     every_n_epochs: 1
   trainer:
-    precision: 32
-    max_epochs: 5 #50
-    min_epochs: 1 #5
+    precision: 16
+    max_epochs: 5
+    min_epochs: 1
diff --git a/docs/tutorials/model_training/ipu_training.ipynb b/docs/tutorials/model_training/ipu_training.ipynb
index c5e1a83c7..cc32c7c97 100644
--- a/docs/tutorials/model_training/ipu_training.ipynb
+++ b/docs/tutorials/model_training/ipu_training.ipynb
@@ -895,6 +895,7 @@
     "logger.info(metrics)\n",
     "\n",
     "predictor = load_predictor(cfg, model_class, model_kwargs, metrics)\n",
+    "predictor.set_max_nodes_edges_per_graph(datamodule, stages=[\"train\", \"val\"])\n",
     "\n",
     "logger.info(predictor.model)\n",
     "logger.info(ModelSummary(predictor, max_depth=4))\n",
@@ -1427,13 +1428,14 @@
    ],
    "source": [
     "ckpt_path = trainer.checkpoint_callbacks[0].best_model_path\n",
+    "predictor.set_max_nodes_edges_per_graph(datamodule, stages=['test'])\n",
     "trainer.test(model=predictor, datamodule=datamodule, ckpt_path=ckpt_path)"
    ]
   }
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3.8.10 ('goli_ipu')",
+   "display_name": "3.1.0+1181_poptorch",
    "language": "python",
    "name": "python3"
   },
@@ -1447,11 +1449,11 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.10"
+   "version": "3.8.10 (default, Jun 22 2022, 20:18:18) \n[GCC 9.4.0]"
   },
   "vscode": {
    "interpreter": {
-    "hash": "d395e6d595d28ecd4e8cd02c8025aead3eecd22bf68b21a59b341f1b0aea0472"
+    "hash": "bc02f5cdc78e740d0e3d50b3e0a92193dd610f8dfa26c6f9f34c1bd27276294a"
    }
   },
   "widgets": {
diff --git a/docs/tutorials/model_training/ipu_training_demo.py b/docs/tutorials/model_training/ipu_training_demo.py
new file mode 100644
index 000000000..a881a0705
--- /dev/null
+++ b/docs/tutorials/model_training/ipu_training_demo.py
@@ -0,0 +1,57 @@
+# General imports
+import os
+from os.path import dirname, abspath
+import yaml
+from copy import deepcopy
+from omegaconf import DictConfig
+import timeit
+from loguru import logger
+from pytorch_lightning.utilities.model_summary import ModelSummary
+
+# Current project imports
+import goli
+from goli.config._loader import load_datamodule, load_metrics, load_architecture, load_predictor, load_trainer
+from goli.utils.safe_run import SafeRun
+
+# from torch_geometric.nn.aggr import Aggregation
+# Aggregation.set_validate_args(False)
+
+# WandB
+import wandb
+
+# Set up the working directory
+MAIN_DIR = dirname(dirname(abspath(goli.__file__)))
+# CONFIG_FILE= "expts/configs/config_ipu_allsizes.yaml"
+CONFIG_FILE = "docs/tutorials/model_training/config_ipu_tutorials.yaml"
+os.chdir(MAIN_DIR)
+
+with open(os.path.join(MAIN_DIR, CONFIG_FILE), "r") as f:
+    cfg = yaml.safe_load(f)
+
+# Load and initialize the dataset
+datamodule = load_datamodule(cfg)
+datamodule.prepare_data()
+
+# Initialize the network
+model_class, model_kwargs = load_architecture(
+    cfg,
+    in_dims=datamodule.in_dims,
+)
+
+metrics = load_metrics(cfg)
+logger.info(metrics)
+
+predictor = load_predictor(cfg, model_class, model_kwargs, metrics)
+predictor.set_max_nodes_edges_per_graph(datamodule, stages=["train", "val"])
+
+logger.info(predictor.model)
+logger.info(ModelSummary(predictor, max_depth=4))
+
+trainer = load_trainer(cfg, "tutorial-run")
+
+# Run the model training
+with SafeRun(name="TRAINING", raise_error=cfg["constants"]["raise_train_error"], verbose=True):
+    trainer.fit(model=predictor, datamodule=datamodule)
+
+# Exit WandB
+wandb.finish()
diff --git a/env.yml b/env.yml
index db02c0ec0..b80a95d4e 100644
--- a/env.yml
+++ b/env.yml
@@ -30,8 +30,8 @@ dependencies:
   - cudatoolkit  # works also with CPU-only system.
   - pytorch >=1.10.2
   - tensorboard
-  - pytorch-lightning >=1.5
-  - torchmetrics >=0.7.0
+  - pytorch-lightning >=1.5,<1.9
+  - torchmetrics >=0.7.0,<0.11
   - dgl >=0.5.2
   - hydra-core >=1.0
   - ogb
diff --git a/expts/configs/config_gpu_qm9.yaml b/expts/configs/config_gpu_qm9.yaml
index c69da5ad1..e771b88c0 100644
--- a/expts/configs/config_gpu_qm9.yaml
+++ b/expts/configs/config_gpu_qm9.yaml
@@ -221,7 +221,7 @@ predictor:
     cv: mse_ipu
   random_seed: *seed
   optim_kwargs:
-    lr: 1.e-3
+    lr: 1.e-4 #1.e-3
     # weight_decay: 1.e-7
   torch_scheduler_kwargs:
     #module_type: ReduceLROnPlateau
diff --git a/expts/configs/ipu.config b/expts/configs/ipu.config
index 3cae2c96a..95013ec3a 100644
--- a/expts/configs/ipu.config
+++ b/expts/configs/ipu.config
@@ -1,5 +1,4 @@
-deviceIterations(4)
+deviceIterations(1)
 replicationFactor(1)
 enableProfiling("graph_analyser")       # The folder where the profile will be stored
-Jit.traceModel(True)
 enableExecutableCaching("pop_compiler_cache")
diff --git a/expts/main_run_get_fingerprints.py b/expts/main_run_get_fingerprints.py
index b5d77b709..5f61992dd 100644
--- a/expts/main_run_get_fingerprints.py
+++ b/expts/main_run_get_fingerprints.py
@@ -25,7 +25,6 @@
 
 
 def main() -> None:
-
     LIST_CONCAT_LAST_LAYERS = [1, 0, [1, 2], [0, 1, 2]]
     DATA_NAME_ALL = ["molbace"]  # , "mollipo", "moltox21", "molHIV"]
     MODEL_PATH = "gs://goli-public/pretrained-models"
diff --git a/expts/main_run_multitask.py b/expts/main_run_multitask.py
index edc0eb5ec..4a4b5c251 100644
--- a/expts/main_run_multitask.py
+++ b/expts/main_run_multitask.py
@@ -28,7 +28,8 @@
 # Set up the working directory
 MAIN_DIR = dirname(dirname(abspath(goli.__file__)))
 # CONFIG_FILE = "expts/configs/config_molPCQM4Mv2.yaml"
-CONFIG_FILE = "expts/configs/config_ipu_qm9.yaml"
+# CONFIG_FILE = "expts/configs/config_ipu_qm9.yaml"
+CONFIG_FILE = "expts/configs/config_gpu_qm9.yaml"
 # CONFIG_FILE = "expts/configs/config_gpu_qm9.yaml"
 # CONFIG_FILE = "expts/configs/config_ipu_reproduce.yaml"
 os.chdir(MAIN_DIR)
@@ -63,10 +64,18 @@ def main(cfg: DictConfig, run_name: str = "main", add_date_time: bool = True) ->
     save_params_to_wandb(trainer.logger, cfg, predictor, datamodule)
 
     datamodule.prepare_data()
+
+    # Determine the max num nodes and edges in training and validation
+    predictor.set_max_nodes_edges_per_graph(datamodule, stages=["train", "val"])
+
     # Run the model training
     with SafeRun(name="TRAINING", raise_error=cfg["constants"]["raise_train_error"], verbose=True):
         trainer.fit(model=predictor, datamodule=datamodule)
 
+    # Determine the max num nodes and edges in testing
+    predictor.set_max_nodes_edges_per_graph(datamodule, stages=["test"])
+
+    # Run the model testing
     with SafeRun(name="TESTING", raise_error=cfg["constants"]["raise_train_error"], verbose=True):
         trainer.test(model=predictor, datamodule=datamodule)  # , ckpt_path=ckpt_path)
 
diff --git a/expts/main_run_predict.py b/expts/main_run_predict.py
index 89be90a09..3006d0074 100644
--- a/expts/main_run_predict.py
+++ b/expts/main_run_predict.py
@@ -29,7 +29,6 @@
 
 
 def main(cfg: DictConfig) -> None:
-
     cfg = deepcopy(cfg)
 
     # Load and initialize the dataset
diff --git a/goli/cli/data.py b/goli/cli/data.py
index cfd58d5c1..3a7e51ffa 100644
--- a/goli/cli/data.py
+++ b/goli/cli/data.py
@@ -36,7 +36,6 @@ def data_cli():
     help="Whether to extract the dataset if it's a zip file.",
 )
 def download(name, output, progress):
-
     args = {}
     args["name"] = name
     args["output_path"] = output
@@ -52,6 +51,5 @@ def download(name, output, progress):
 
 @data_cli.command(name="list", help="List available Goli dataset.")
 def list():
-
     logger.info("Goli datasets:")
     logger.info(goli.data.utils.list_goli_datasets())
diff --git a/goli/config/_loader.py b/goli/config/_loader.py
index c301f0e53..e183cb439 100644
--- a/goli/config/_loader.py
+++ b/goli/config/_loader.py
@@ -1,4 +1,4 @@
-from typing import Callable, Dict, Mapping, Type, Union, Any
+from typing import Callable, Dict, Mapping, Type, Union, Any, Optional, List, Tuple
 
 # Misc
 import os
@@ -24,9 +24,12 @@
 from goli.nn.architectures import FullGraphNetwork, FullGraphMultiTaskNetwork
 from goli.trainer.predictor import PredictorModule
 from goli.utils.spaces import DATAMODULE_DICT
-from goli.ipu.ipu_wrapper import PredictorModuleIPU, IPUPluginGoli
 from goli.ipu.ipu_utils import import_poptorch, load_ipu_options
 from goli.data.datamodule import MultitaskFromSmilesDataModule
+from goli.data.datamodule import BaseDataModule
+
+# Weights and Biases
+from pytorch_lightning import Trainer
 
 
 def get_accelerator(
@@ -73,7 +76,7 @@ def get_accelerator(
     return acc_type
 
 
-def load_datamodule(config: Union[omegaconf.DictConfig, Dict[str, Any]]) -> "goli.datamodule.BaseDataModule":
+def load_datamodule(config: Union[omegaconf.DictConfig, Dict[str, Any]]) -> BaseDataModule:
     """
     Load the datamodule from the specified configurations at the key
     `datamodule: args`.
@@ -250,6 +253,8 @@ def load_predictor(
     """
 
     if get_accelerator(config) == "ipu":
+        from goli.ipu.ipu_wrapper import PredictorModuleIPU
+
         predictor_class = PredictorModuleIPU
     else:
         predictor_class = PredictorModule
@@ -263,7 +268,7 @@ def load_predictor(
     )
 
     # mup base shapes
-    mup_base_path = config["architecture"].pop("mup_base_path")
+    mup_base_path = config["architecture"].pop("mup_base_path", None)
     predictor = load_mup(mup_base_path, predictor)
 
     return predictor
@@ -299,7 +304,7 @@ def load_trainer(config: Union[omegaconf.DictConfig, Dict[str, Any]], run_name:
     cfg_trainer = deepcopy(config["trainer"])
 
     # Define the IPU plugin if required
-    plugins = []
+    strategy = None
     accelerator = get_accelerator(config)
     ipu_file = "expts/configs/ipu.config"
     if accelerator == "ipu":
@@ -309,7 +314,9 @@ def load_trainer(config: Union[omegaconf.DictConfig, Dict[str, Any]], run_name:
             model_name=config["constants"]["name"],
             gradient_accumulation=config["trainer"]["trainer"].get("accumulate_grad_batches", None),
         )
-        plugins = IPUPluginGoli(training_opts=training_opts, inference_opts=inference_opts)
+        from goli.ipu.ipu_wrapper import DictIPUStrategy
+
+        strategy = DictIPUStrategy(training_opts=training_opts, inference_opts=inference_opts)
 
     # Set the number of gpus to 0 if no GPU is available
     _ = cfg_trainer["trainer"].pop("accelerator", None)
@@ -347,7 +354,7 @@ def load_trainer(config: Union[omegaconf.DictConfig, Dict[str, Any]], run_name:
 
     trainer = Trainer(
         detect_anomaly=True,
-        plugins=plugins,
+        strategy=strategy,
         accelerator=accelerator,
         ipus=ipus,
         gpus=gpus,
diff --git a/goli/config/zinc_default_multitask_fulldgl.yaml b/goli/config/zinc_default_multitask_fulldgl.yaml
index d978c8098..0301f24c6 100644
--- a/goli/config/zinc_default_multitask_fulldgl.yaml
+++ b/goli/config/zinc_default_multitask_fulldgl.yaml
@@ -112,7 +112,7 @@ architecture:     # The parameters for the full graph network are taken from `co
     residual_type: none
 
   task_heads:     # Set as null to avoid task heads. Recall that the arguments for the TaskHeads is a List of TaskHeadParams
-    - task_name: "task_1"
+    task_1:
       out_dim: 5
       hidden_dims: [5, 6, 7]
       #depth: none                          # Not needed if we have hidden_dims
@@ -121,7 +121,7 @@ architecture:     # The parameters for the full graph network are taken from `co
       dropout: 0.2
       normalization: none
       residual_type: none
-    - task_name: "task_2"
+    task_2:
       out_dim: 3
       hidden_dims: [8, 9, 10]
       activation: relu
@@ -129,7 +129,7 @@ architecture:     # The parameters for the full graph network are taken from `co
       dropout: 0.2
       normalization: none
       residual_type: none
-    - task_name: "task_3"
+    task_3:
       out_dim: 4
       hidden_dims: [2, 2, 2]
       activation: relu
diff --git a/goli/data/datamodule.py b/goli/data/datamodule.py
index 1b6262e54..43c96b9e5 100644
--- a/goli/data/datamodule.py
+++ b/goli/data/datamodule.py
@@ -423,7 +423,6 @@ def __init__(
 
         self.collate_fn = self.get_collate_fn(collate_fn)
 
-        self.dataset = None
         self.train_ds = None
         self.val_ds = None
         self.test_ds = None
@@ -442,6 +441,7 @@ def train_dataloader(self, **kwargs):
             dataset=self.train_ds,  # type: ignore
             shuffle=True,
             stage=RunningStage.TRAINING,
+            **kwargs,
         )
 
     def val_dataloader(self, **kwargs):
@@ -449,22 +449,23 @@ def val_dataloader(self, **kwargs):
             dataset=self.val_ds,  # type: ignore
             shuffle=False,
             stage=RunningStage.VALIDATING,
+            **kwargs,
         )
 
     def test_dataloader(self, **kwargs):
-
         return self.get_dataloader(
             dataset=self.test_ds,  # type: ignore
             shuffle=False,
             stage=RunningStage.TESTING,
+            **kwargs,
         )
 
     def predict_dataloader(self, **kwargs):
-
         return self.get_dataloader(
             dataset=self.predict_ds,  # type: ignore
             shuffle=False,
             stage=RunningStage.PREDICTING,
+            **kwargs,
         )
 
     @staticmethod
@@ -603,6 +604,84 @@ def _dataloader(self, dataset: Dataset, **kwargs) -> DataLoader:
 
         return loader
 
+    def get_max_num_nodes_datamodule(self, stages: Optional[List[str]] = None) -> int:
+        """
+        Get the maximum number of nodes across all datasets from the datamodule
+
+        Parameters:
+            datamodule: The datamodule from which to extract the maximum number of nodes
+            stages: The stages from which to extract the max num nodes.
+                Possible values are ["train", "val", "test", "predict"].
+                If None, all stages are considered.
+
+        Returns:
+            max_num_nodes: The maximum number of nodes across all datasets from the datamodule
+        """
+
+        allowed_stages = ["train", "val", "test", "predict"]
+        if stages is None:
+            stages = allowed_stages
+        for stage in stages:
+            assert stage in allowed_stages, f"stage value `{stage}` not allowed."
+
+        max_num_nodes = 0
+        # Max number of nodes in the training dataset
+        if (self.train_ds is not None) and ("train" in stages):
+            max_num_nodes = max(max_num_nodes, self.train_ds.max_num_nodes_per_graph)
+
+        # Max number of nodes in the validation dataset
+        if (self.val_ds is not None) and ("val" in stages):
+            max_num_nodes = max(max_num_nodes, self.val_ds.max_num_nodes_per_graph)
+
+        # Max number of nodes in the test dataset
+        if (self.test_ds is not None) and ("test" in stages):
+            max_num_nodes = max(max_num_nodes, self.test_ds.max_num_nodes_per_graph)
+
+        # Max number of nodes in the predict dataset
+        if (self.predict_ds is not None) and ("predict" in stages):
+            max_num_nodes = max(max_num_nodes, self.predict_ds.max_num_nodes_per_graph)
+
+        return max_num_nodes
+
+    def get_max_num_edges_datamodule(self, stages: Optional[List[str]] = None) -> int:
+        """
+        Get the maximum number of edges across all datasets from the datamodule
+
+        Parameters:
+            datamodule: The datamodule from which to extract the maximum number of nodes
+            stages: The stages from which to extract the max num nodes.
+                Possible values are ["train", "val", "test", "predict"].
+                If None, all stages are considered.
+
+        Returns:
+            max_num_edges: The maximum number of edges across all datasets from the datamodule
+        """
+
+        allowed_stages = ["train", "val", "test", "predict"]
+        if stages is None:
+            stages = allowed_stages
+        for stage in stages:
+            assert stage in allowed_stages, f"stage value `{stage}` not allowed."
+
+        max_num_edges = 0
+        # Max number of nodes/edges in the training dataset
+        if (self.train_ds is not None) and ("train" in stages):
+            max_num_edges = max(max_num_edges, self.train_ds.max_num_edges_per_graph)
+
+        # Max number of nodes/edges in the validation dataset
+        if (self.val_ds is not None) and ("val" in stages):
+            max_num_edges = max(max_num_edges, self.val_ds.max_num_edges_per_graph)
+
+        # Max number of nodes/edges in the test dataset
+        if (self.test_ds is not None) and ("test" in stages):
+            max_num_edges = max(max_num_edges, self.test_ds.max_num_edges_per_graph)
+
+        # Max number of nodes/edges in the predict dataset
+        if (self.predict_ds is not None) and ("predict" in stages):
+            max_num_edges = max(max_num_edges, self.predict_ds.max_num_edges_per_graph)
+
+        return max_num_edges
+
 
 class DatasetProcessingParams:
     def __init__(
@@ -815,7 +894,6 @@ def __init__(
         self.val_singletask_datasets = None
         self.test_singletask_datasets = None
 
-        self.dataset = None
         self.train_ds = None
         self.val_ds = None
         self.test_ds = None
@@ -1617,7 +1695,7 @@ def __len__(self) -> int:
         num_elements = 0
         for task, args in self.task_dataset_processing_params.items():
             if args.df is None:
-                df = self._read_csv(args.df_path, usecols=args.smiles_col)
+                df = self._read_csv(args.df_path, usecols=[args.smiles_col])
                 num_elements += len(df)
             else:
                 num_elements += len(args.df)
@@ -1756,7 +1834,6 @@ def _load_dataset(self, metadata: dict, sample_size: Optional[int] = None):
         dataset_dir = base_dir / metadata["download_name"]
 
         if not dataset_dir.exists():
-
             # Create cache filepath for zip file and associated folder
             dataset_path = base_dir / f"{metadata['download_name']}.zip"
 
diff --git a/goli/data/utils.py b/goli/data/utils.py
index 497f7d6c1..ffd5d8613 100644
--- a/goli/data/utils.py
+++ b/goli/data/utils.py
@@ -58,7 +58,6 @@ def download_goli_dataset(name: str, output_path: str, extract_zip: bool = True,
         goli.utils.fs.copy(dataset_path_source, dataset_path_destination, progress=progress)
 
         if extract_zip and str(dataset_path_destination).endswith(".zip"):
-
             # Unzip the dataset
             with zipfile.ZipFile(dataset_path_destination, "r") as zf:
                 zf.extractall(output_path)
diff --git a/goli/features/featurizer.py b/goli/features/featurizer.py
index eab541049..808a444fd 100644
--- a/goli/features/featurizer.py
+++ b/goli/features/featurizer.py
@@ -33,7 +33,6 @@ def to_dense_array(array, dtype=None):
 
 
 def _mask_nans_inf(mask_nan, array, array_name):
-
     if (mask_nan is None) or (array is None):
         return array
 
@@ -108,13 +107,11 @@ def get_mol_atomic_features_onehot(mol: dm.Mol, property_list: List[str]) -> Dic
     prop_dict = {}
 
     for prop in property_list:
-
         prop = prop.lower()
         prop_name = prop
 
         property_array = []
         for ii, atom in enumerate(mol.GetAtoms()):
-
             if prop in ["atomic-number"]:
                 one_hot = one_of_k_encoding(atom.GetSymbol(), nmp.ATOM_LIST)
             elif prop in ["degree"]:
@@ -237,16 +234,13 @@ def get_mol_atomic_features_float(
     atom_list = list(mol.GetAtoms())
 
     for prop in property_list:
-
         prop_name = None
 
         property_array = np.zeros(mol.GetNumAtoms(), dtype=np.float16)
         for ii, atom in enumerate(atom_list):
-
             val = None
 
             if isinstance(prop, str):
-
                 prop = prop.lower()
                 prop_name = prop
 
diff --git a/goli/features/nmp.py b/goli/features/nmp.py
index fbeb51ce7..8557c5921 100644
--- a/goli/features/nmp.py
+++ b/goli/features/nmp.py
@@ -10,6 +10,7 @@
     PERIODIC_TABLE = pd.read_csv(f)
 PERIODIC_TABLE = PERIODIC_TABLE.set_index("AtomicNumber")
 
+
 # Small function to convert strings to floats
 def float_or_none(string):
     try:
diff --git a/goli/features/spectral.py b/goli/features/spectral.py
index be7c8d31a..5f56b5af8 100644
--- a/goli/features/spectral.py
+++ b/goli/features/spectral.py
@@ -16,7 +16,6 @@ def compute_laplacian_positional_eigvecs(
     disconnected_comp: bool = True,
     normalization: str = "none",
 ) -> Tuple[np.ndarray, np.ndarray]:
-
     # Sparsify the adjacency patrix
     if issparse(adj):
         adj = adj.astype(np.float64)
@@ -57,7 +56,6 @@ def compute_laplacian_positional_eigvecs(
 
 
 def _get_positional_eigvecs(matrix, num_pos: int):
-
     mat_len = matrix.shape[0]
     eigvals, eigvecs = eig(matrix.todense())
 
diff --git a/goli/ipu/ipu_dataloader.py b/goli/ipu/ipu_dataloader.py
index 85bee4420..699c20366 100644
--- a/goli/ipu/ipu_dataloader.py
+++ b/goli/ipu/ipu_dataloader.py
@@ -5,6 +5,7 @@
 from dataclasses import dataclass
 import numpy as np
 from loguru import logger
+from torch import Tensor
 
 import torch
 from torch_geometric.data import Data, Batch, Dataset
@@ -34,7 +35,6 @@ class IPUDataloaderOptions:
     mode: "poptorch.DataLoaderMode" = "Sync"
 
     def set_kwargs(self):
-
         # Get the maximum number of nodes
         if self.max_num_nodes is not None:
             assert (
@@ -139,7 +139,7 @@ def __call__(self, batch: Batch) -> Batch:
                 dataset_max_edges_per_graph=self.dataset_max_edges_per_graph,
             )
 
-            local_batch["features"], local_batch["_types_conversion"] = transform(local_batch["features"])
+            local_batch["features"] = transform(local_batch["features"])
             all_batches.append(local_batch)
 
         out_batch = {}
@@ -157,12 +157,17 @@ def __call__(self, batch: Batch) -> Batch:
 
         # TODO: Make this more robust, instead of hard-coding the keys
         out_batch["features"] = stacked_features
-        out_batch["_types_conversion"] = [this_batch["_types_conversion"] for this_batch in all_batches]
-        out_batch["_batch_idx"] = torch.as_tensor(range(len(all_batches)), dtype=torch.int64).unsqueeze(-1)
+        out_batch["_batch_idx"] = torch.as_tensor(range(len(all_batches)), dtype=torch.int32).unsqueeze(-1)
         for key in all_batches[0].keys():
             if key not in ("features", "labels", "_types_conversion"):
                 out_batch[key] = [this_batch[key] for this_batch in all_batches]
 
+        for data_key, data_val in out_batch.items():
+            if isinstance(data_val, Batch):
+                for sub_key, sub_val in data_val.items():
+                    if isinstance(sub_val, Tensor) and sub_val.dtype == torch.int64:
+                        out_batch[data_key][sub_key] = sub_val.to(torch.int32)
+
         return out_batch
 
 
@@ -372,33 +377,7 @@ def __call__(self, batch: Batch) -> Batch:
         if "num_nodes" in new_batch:
             new_batch.num_nodes = self.max_num_nodes
 
-        new_batch.dataset_max_nodes_per_graph = torch.as_tensor(
-            [self.dataset_max_nodes_per_graph], dtype=torch.int32
-        )
-        new_batch.dataset_max_edges_per_graph = torch.as_tensor(
-            [self.dataset_max_edges_per_graph], dtype=torch.int32
-        )
-
-        # Convert integer types to smaller integers for faster transfer of data to IPUs
-        _types_conversion = {}
-        int_types = [torch.int16, torch.int32, torch.int64]
-        for key, val in new_batch.items():
-            if isinstance(val, torch.Tensor) and (val.dtype in int_types[1:]):
-                for this_int_type in int_types:
-                    if (val.max() <= torch.iinfo(this_int_type).max) and (
-                        val.min() >= torch.iinfo(this_int_type).min
-                    ):
-                        _types_conversion[key] = val.dtype
-                        new_batch[key] = val.to(this_int_type)
-                        break
-
-        # Convert float types to smaller floats for faster transfer of data to IPUs
-        float_types = [torch.float16, torch.float32, torch.float64]
-        for key, val in new_batch.items():
-            if isinstance(val, torch.Tensor) and (val.dtype in float_types[1:]):
-                new_batch[key] = val.to(torch.float16)
-
-        return new_batch, _types_conversion
+        return new_batch
 
     def __repr__(self) -> str:
         s = f"{self.__class__.__name__}("
diff --git a/goli/ipu/ipu_losses.py b/goli/ipu/ipu_losses.py
index 63513a11d..2ba2f39c3 100644
--- a/goli/ipu/ipu_losses.py
+++ b/goli/ipu/ipu_losses.py
@@ -53,7 +53,6 @@ class MSELossIPU(MSELoss):
     """
 
     def forward(self, input: Tensor, target: Tensor) -> Tensor:
-
         target = target.clone()
         input = input.clone()
 
@@ -78,7 +77,6 @@ class L1LossIPU(L1Loss):
     """
 
     def forward(self, input: Tensor, target: Tensor) -> Tensor:
-
         target = target.clone()
         input = input.clone()
 
diff --git a/goli/ipu/ipu_simple_lightning.py b/goli/ipu/ipu_simple_lightning.py
index 51a0bb508..77333e3da 100644
--- a/goli/ipu/ipu_simple_lightning.py
+++ b/goli/ipu/ipu_simple_lightning.py
@@ -1,6 +1,6 @@
 # Copyright (c) 2021 Graphcore Ltd. All rights reserved.
 import pytorch_lightning as pl
-from pytorch_lightning.plugins import IPUPlugin
+from pytorch_lightning.strategies import IPUStrategy
 from pytorch_lightning.loggers import WandbLogger
 
 import torch
@@ -84,7 +84,7 @@ def validation_step(self, batch, _):
 
     # PopTorch doesn't currently support logging within steps. Use the Lightning
     # callback hooks instead.
-    def on_train_batch_end(self, outputs, batch, batch_idx, dataloader_idx):
+    def on_train_batch_end(self, outputs, batch, batch_idx):
         self.log("StepLoss", outputs["loss"])
 
     def validation_epoch_end(self, outputs):
@@ -107,7 +107,6 @@ def configure_optimizers(self):
 
 
 if __name__ == "__main__":
-
     torch.manual_seed(SEED)
 
     # Create the model as usual.
@@ -137,21 +136,18 @@ def configure_optimizers(self):
         import poptorch
 
         training_opts = poptorch.Options()
-        training_opts.Jit.traceModel(True)
         inference_opts = poptorch.Options()
-        inference_opts.Jit.traceModel(True)
 
         # Set the seeds
         training_opts.randomSeed(SEED)
         inference_opts.randomSeed(SEED)
         ipus = 1
-        plugins = IPUPlugin(training_opts=training_opts, inference_opts=inference_opts)
+        strategy = IPUStrategy(training_opts=training_opts, inference_opts=inference_opts)
 
     trainer = pl.Trainer(
         logger=WandbLogger(),
         ipus=ipus,
         max_epochs=3,
-        progress_bar_refresh_rate=20,
         log_every_n_steps=1,
         plugins=plugins,
     )
diff --git a/goli/ipu/ipu_wrapper.py b/goli/ipu/ipu_wrapper.py
index d2094c817..652efa8ea 100644
--- a/goli/ipu/ipu_wrapper.py
+++ b/goli/ipu/ipu_wrapper.py
@@ -1,14 +1,21 @@
-from typing import Dict, Any, Optional, Callable, Union, Type, Tuple
+from typing import Dict, Any, Optional, Callable, Union, Type, Tuple, Iterable
 
 from torch_geometric.data import Batch
 from torch import Tensor
-from inspect import _ParameterKind
-from pytorch_lightning.plugins import IPUPlugin
+from pytorch_lightning.strategies import IPUStrategy
+from pytorch_lightning.utilities.types import STEP_OUTPUT
 from pytorch_lightning.trainer.states import RunningStage
 
 from goli.trainer.predictor import PredictorModule
 from goli.ipu.ipu_utils import import_poptorch
 
+import torch
+from torch_geometric.data import Data, Batch
+from torch_geometric.data.data import BaseData
+from loguru import logger
+
+poptorch = import_poptorch()
+
 
 def remove_pad_loss(preds: Dict[str, Tensor], targets: Dict[str, Tensor]):
     """
@@ -23,84 +30,73 @@ def remove_pad_loss(preds: Dict[str, Tensor], targets: Dict[str, Tensor]):
     return preds
 
 
-class IPUPluginGoli(IPUPlugin):
+class DictIPUStrategy(IPUStrategy):
+    def _step(self, stage: RunningStage, *args: Any, **kwargs: Any) -> STEP_OUTPUT:
+        args = self._prepare_input(args)
+        args = args[0]
+        poptorch_model = self.poptorch_models[stage]
+        self.lightning_module._running_torchscript = True
+        for key_to_drop in ["_batch_idx", "mol_ids", "smiles"]:
+            args.pop(key_to_drop)
+        out = poptorch_model(**args)
+        self.lightning_module._running_torchscript = False
+        return out
+
+
+class PyGArgsParser(poptorch.ICustomArgParser):
     """
-    `IPUPluginGoli` modifies the `IPUPlugin` for compatibility with the Goli and Pytorch-Lightning training pipeline.
-    Modifies the `.self` method to make it compatible with the rest of Goli.
+    This class is responsible for converting a PyG Batch from and to
+    a tensor of tuples. This allows PyG Batch to be used as inputs to
+    IPU programs. Copied from poppyg repo, in the future import from
+    the repo directly.
     """
 
-    def _step(self, stage: RunningStage, *args: Any, **kwargs: Any):
+    @staticmethod
+    def sortedTensorKeys(struct: BaseData) -> Iterable[str]:
         """
-        The `Goli` Dataloader generates Dictionaries as inputs for the `Predictor`.
-        But the IPU does not support the following:
-          - Dictionaries of inputs. Must pass Tuples. NamedTuples won't work, they will be converted to Tuples.
-          - Only Tensors can be passed to the model. Anything else (strings, numpy array, etc) will cause compilation to fail.
+        Find all the keys that map to a tensor value in struct. The keys
+        are returned in sorted order.
+        """
+        all_keys = sorted(struct.keys)
 
-        This function overloads the regular `_step` function to do the following:
-          1. Create `Tuple[Tensor]` from the inputs of type `Dict[Tensor]` or `Tensor` or `pyg.Data.Batch` and pass them to `super()._step`
-          2. Save the keys of the `Dict` structure as attributes of the model: `self.model.module._keys_tensor_dict` and `self.model.module._keys_tensor`
-          3. Convert every non-Tensor to Tuples and pass to the module directly.
-          4. Run the `super()._step` function, then remove every parameter that was added to the model
+        def isTensor(k: str) -> bool:
+            return isinstance(struct[k], torch.Tensor)
 
-        The model must take care of re-building the Dict[Tensor] and `Batch` from the tuples. See class `PredictorModuleIPU`.
+        return filter(isTensor, all_keys)
 
+    def yieldTensors(self, struct: BaseData):
+        """
+        yield every torch.Tensor in struct in sorted order
         """
+        for k in self.sortedTensorKeys(struct):
+            yield struct[k]
 
-        # Arguments for the loop
-        new_args, all_keys = [], []
-        keys_tensor_dict, keys_batch, keys_tensor, keys_others = {}, {}, {}, {}
-
-        # Loop every argument. If a dict or pyg graph is found, split into tensors
-        for data_key, data_val in args[0].items():
-            if isinstance(data_val, (Dict, Batch)):
-                for sub_key, sub_val in data_val.items():
-                    if isinstance(sub_val, Tensor):
-                        new_args.append(sub_val)
-                        all_keys.append(sub_key)
-
-                        # Append the keys for the tensors
-                        if isinstance(data_val, Dict):
-                            if data_key not in keys_tensor_dict:
-                                keys_tensor_dict[data_key] = {}
-                            keys_tensor_dict[data_key][sub_key] = len(all_keys) - 1
-
-                        # Append the keys for the pyg Batch
-                        elif isinstance(data_val, Batch):
-                            if data_key not in keys_batch:
-                                keys_batch[data_key] = {}
-                            keys_batch[data_key][sub_key] = len(all_keys) - 1
-                    else:
-                        if data_key not in keys_others:
-                            keys_others[data_key] = {}
-                        keys_others[data_key][sub_key] = sub_val
-
-            elif isinstance(data_val, Tensor):
-                new_args.append(data_val)
-                all_keys.append(data_key)
-                keys_tensor[data_key] = len(all_keys) - 1
-            else:
-                keys_others[data_key] = data_val
-
-        # Tell the module what are the labels associated to the labels and batches
-        self.model.module._keys_tensor_dict = keys_tensor_dict
-        self.model.module._keys_tensor = keys_tensor
-        self.model.module._keys_batch = keys_batch
-        self.model.module._keys_others = keys_others
-
-        # Walk-around to set the variable names for the poptorch model
-        self.poptorch_models[stage]._args_parser._varnames = all_keys
-        self.poptorch_models[stage]._args_parser._var_kinds = [_ParameterKind.VAR_POSITIONAL] * len(all_keys)
-
-        # Run the step using only tuple of tensors
-        out = super()._step(stage, *new_args, **kwargs)
-
-        # Remove the keys from the module after the step is executed
-        self.model.module._keys_tensor_dict = None
-        self.model.module._keys_tensor = None
-        self.model.module._keys_batch = None
-        self.model.module._keys_others = None
+    def reconstruct(self, original_structure: BaseData, tensor_iterator: Iterable[Tensor]):
+        """
+        Create a new instance with the same class type as the
+        original_structure. This new instance will be initialized with tensors
+        from the provided iterator and uses the same sorted keys from the
+        yieldTensors() implementation.
+        """
+        tensor_keys = self.sortedTensorKeys(original_structure)
+        kwargs = {k: next(tensor_iterator) for k in tensor_keys}
 
-        return out
+        for k in original_structure.keys:
+            if k not in kwargs:
+                # copy non-tensor properties to the new instance
+                kwargs[k] = original_structure[k]
+
+        cls = original_structure.__class__
+
+        if issubclass(cls, Batch):
+            kwargs["_base_cls"] = Data
+            return Batch(**kwargs)
+
+        return cls(**kwargs)
+
+
+# PyG uses the BaseData object as the root for data and batch objects
+poptorch.registerCustomArgParser(BaseData, PyGArgsParser())
 
 
 class PredictorModuleIPU(PredictorModule):
@@ -128,198 +124,52 @@ def compute_loss(
             preds, targets, weights, loss_fun, target_nan_mask, multitask_handling
         )
 
-    def on_train_batch_end(self, outputs, batch, batch_idx, dataloader_idx):
-        outputs = {"loss/train": outputs["loss"].mean()}
-        super().on_train_batch_end(outputs, batch, batch_idx, dataloader_idx)
+    def on_train_batch_end(self, outputs, batch, batch_idx):
+        outputs["loss/train"] = outputs["loss"].mean()
+        super().on_train_batch_end(outputs, batch, batch_idx)
 
-    def training_step(self, *inputs) -> Dict[str, Any]:
-        # Build a dictionary from the tuples
-        dict_input = self._build_dict_input(*inputs)
-        concatenated_metrics_logs = super().training_step(dict_input, to_cpu=False)
-        loss = concatenated_metrics_logs.pop("loss")
-        loss = self.poptorch.identity_loss(loss, reduction="mean")
-        return loss  # Limitation that only the loss can be returned
+    def training_step(self, features, labels) -> Dict[str, Any]:
+        features, labels = self.squeeze_input_dims(features, labels)
+        dict_input = {"features": features, "labels": labels}
+        step_dict = super().training_step(dict_input, to_cpu=False)
 
-    def validation_step(self, *inputs) -> Dict[str, Any]:
-        # Build a dictionary from the tuples
-        dict_input = self._build_dict_input(*inputs)
+        loss = step_dict.pop("loss")
+        step_dict["loss"] = self.poptorch.identity_loss(loss, reduction="mean")
+        return step_dict
+
+    def validation_step(self, features, labels) -> Dict[str, Any]:
+        features, labels = self.squeeze_input_dims(features, labels)
+        dict_input = {"features": features, "labels": labels}
         step_dict = super().validation_step(dict_input, to_cpu=False)
 
-        # The output dict must be converted to a tuple
-        step_dict = self._clean_output_batch(step_dict)
         return step_dict
 
-    def test_step(self, *inputs) -> Dict[str, Any]:
+    def test_step(self, features, labels) -> Dict[str, Any]:
         # Build a dictionary from the tuples
-        dict_input = self._build_dict_input(*inputs)
+        features, labels = self.squeeze_input_dims(features, labels)
+        dict_input = {"features": features, "labels": labels}
         step_dict = super().test_step(dict_input, to_cpu=False)
 
-        # The output dict must be converted to a tuple
-        step_dict = self._clean_output_batch(step_dict)
         return step_dict
 
-    def predict_step(self, *inputs) -> Dict[str, Any]:
+    def predict_step(self, **inputs) -> Dict[str, Any]:
         # Build a dictionary from the tuples
-        dict_input = self._build_dict_input(*inputs)
+        dict_input = inputs
         step_dict = super().predict_step(dict_input, to_cpu=False)
 
-        # The output dict must be converted to a tuple
-        step_dict = self._clean_output_batch(step_dict)
         return step_dict
 
-    def training_epoch_end(self, outputs: Dict[str, Any]):
-        # Limited. Since only the loss can be returned in `training_step`
-        return
-
-    def validation_epoch_end(self, outputs: Dict[str, Any]):
-        # Retrieve the dict structure of the output batch from the tuple
-        retrieved_outputs = self._retrieve_output_batch(outputs)
-        return super().validation_epoch_end(retrieved_outputs)
-
-    def test_epoch_end(self, outputs: Dict[str, Any]):
-        # Retrieve the dict structure of the output batch from the tuple
-        retrieved_outputs = self._retrieve_output_batch(outputs)
-        return super().test_epoch_end(retrieved_outputs)
-
-    def predict_epoch_end(self, outputs: Dict[str, Any]):
-        # Retrieve the dict structure of the output batch from the tuple
-        retrieved_outputs = self._retrieve_output_batch(outputs)
-        return super().test_epoch_end(retrieved_outputs)
-
     def configure_optimizers(self, impl=None):
         if impl is None:
             impl = self.poptorch.optim.Adam
         return super().configure_optimizers(impl=impl)
 
-    def _retrieve_output_batch(self, outputs):
-        """
-        A limitation of the IPU is that only Tuples can be returned from
-        `validation_step` and `test_step`.
-
-        Here, we rebuild a dictionary from the tuples so that the code
-        remains compatible with the original `[STAGE]_epoch_end` methods.
-
-        The keys to rebuild the dict must be contained in `self._output_step_structure`
-        as a `Dict[Dict[str]]`, with `self._output_step_structure["_others"]` representing the keys of
-        the non-nested dictionaries.
-
-        """
-
-        new_outputs = []
-        for batch in outputs:
-            new_outputs.append({})
-            # Get the keys of the nested dictionaries
-            for ii, struct in enumerate(self._output_step_structure):
-                new_outputs[-1][struct[0]] = {}
-                for jj, key in enumerate(struct[1:]):
-                    new_outputs[-1][struct[0]][key] = batch[ii][jj]
-
-            # Pop the non-nested dictionaries, and re-add them
-            others = new_outputs[-1].pop("_others", {})
-            new_outputs[-1].update(others)
-        return new_outputs
-
-    def _clean_output_batch(self, out_batch):
-        """
-        The output batch cannot contain `Dict[Tensor]` or `Dict[Dict[Tensor]]`,
-        only nested tuples. And they must all be the same depth.
-
-        This function converts every `Dict[Tensor]` or `Dict[Dict[Tensor]]` to
-        `Tuple[Tuple[Tensor]]` and stores the keys `self._output_step_structure`.
-
-        In the case of non-nested dict, they will be stored under the key `_others`
-        to ensure that everything has the same nesting depth.
-        """
-
-        # Transform Dict[Tensor] into Dict[Dict[Tensor]] by grouping them
-        cleaned_batch, others = {}, {}
-        for key, val in out_batch.items():
-            if isinstance(val, dict):
-                cleaned_batch[key] = val
-            elif isinstance(val, Tensor):
-                others[key] = val
-        if len(others) > 0:
-            cleaned_batch["_others"] = others
-
-        # Save the structure of the dict somewhere
-        output_step_structure = []
-        for key, val in cleaned_batch.items():
-            output_step_structure.append([key])
-            for sub_key, sub_val in val.items():
-                output_step_structure[-1].append(sub_key)
-        self._output_step_structure = output_step_structure
-
-        # Convert Dict[Dict[Tensor]] into Tuple[Tuple[Tensor]]
-        new_dict = {}
-        for key, val in cleaned_batch.items():
-            new_dict[key] = tuple(val.values())
-        cleaned_batch = tuple(new_dict.values())
-
-        return cleaned_batch
-
-    def _build_dict_input(self, *inputs):
-        """
-        The method `IPUPluginGoli._step` converts the `Dict` structure into Tuples
-        to allow the IPU tracer to compile correctly.
-
-        This method rebuilds the `Dict` structure from the saved keys to allow
-        to use the same code as the CPU or GPU.
-
-        It processes the following attributes, which must be set by `IPUPluginGoli._step`:
-          - _keys_batch: The keys for rebuilding the `pyg.Data.Batch`
-          - _keys_tensor_dict: The keys for rebuilding the `Dict[Dict[Tensor]]`
-          - _keys_tensor: The keys for rebuilding the `Dict[Tensor]`
-          - _keys_others: The keys for non-Tensor inputs, which are passed directly to the model.
-        """
-
-        batch = {}
-
-        # Unsqueeze the first dimension that is always 1 here, but is due to the global batch size
-        # (device_iterations * replication_factor * gradient_accumulation)
-        if all([this_input.shape[0] == 1 for this_input in inputs]):
-            inputs = [this_input.squeeze(0) for this_input in inputs]
-
-        # Initialize the batch of pyg objects
-        for key, pyg_elems in self._keys_batch.items():
-            pyg_batch = Batch()
-            for pyg_key, idx in pyg_elems.items():
-                pyg_batch[pyg_key] = inputs[idx]
-            batch[key] = pyg_batch
-
-        # Initialize the dictionaries of tensors, such as the multitask labels
-        for key, this_dict in self._keys_tensor_dict.items():
-            tensor_dict = {}
-            for tensor_key, idx in this_dict.items():
-                tensor_dict[tensor_key] = inputs[idx]
-            batch[key] = tensor_dict
-
-        # Initialize the tensors
-        for key, idx in self._keys_tensor.items():
-            batch[key] = inputs[idx]
-
-        # Initialize the other elements
-        for key, val in self._keys_others.items():
-            if isinstance(val, dict):
-                if key not in batch.keys():
-                    batch[key] = {}
-                # Update the dict or pyg-Batch
-                for sub_key, sub_val in val.items():
-                    batch[sub_key] = sub_val
-            else:
-                batch[key] = val
-
-        # Get the current index for non-tensor elements
-        batch_idx = batch.pop("_batch_idx")
-        batch_idx = batch_idx.squeeze(-1)
-
-        non_tensor_keys = set(self._keys_others.keys()) - (
-            set(self._keys_batch.keys()) | set(self._keys_tensor.keys()) | set(self._keys_tensor_dict.keys())
-        )
-        for key in non_tensor_keys:
-            batch[key] = batch[key][batch_idx]
+    def squeeze_input_dims(self, features, labels):
+        for key, tensor in features:
+            if isinstance(tensor, torch.Tensor):
+                features[key] = features[key].squeeze(0)
 
-        # Convert the tensors to their full dtype (instead of the reduced dtype used to increase data transfer speed)
-        for key, new_dtype in self._keys_others["_types_conversion"][batch_idx].items():
-            batch["features"][key] = batch["features"][key].to(new_dtype)
+        for key in labels:
+            labels[key] = labels[key].squeeze(0)
 
-        return batch
+        return features, labels
diff --git a/goli/ipu/to_dense_batch.py b/goli/ipu/to_dense_batch.py
index 56fd6960a..20c3c4928 100644
--- a/goli/ipu/to_dense_batch.py
+++ b/goli/ipu/to_dense_batch.py
@@ -54,6 +54,10 @@ def to_dense_batch(
         batch = x.new_zeros(x.size(0), dtype=torch.long)
 
     if batch_size is None:
+        assert x.device.type != "ipu", (
+            "When using the IPU the batch size must be "
+            "provided during compilation instead of determined at runtime"
+        )
         batch_size = int(batch.max()) + 1
 
     num_nodes = scatter_add(batch.new_ones(x.size(0)), batch, dim=0, dim_size=batch_size)
@@ -77,16 +81,24 @@ def to_dense_batch(
         idx[idx >= size[0]] = size[0] - 1
 
     # Raise error if num_nodes > max_num_nodes
-    assert (
-        num_nodes <= max_num_nodes_per_graph
-    ).all(), f"Encountered graphs with {num_nodes.max()} nodes, greater than `max_num_nodes = {max_num_nodes_per_graph}`"
-
-    ##### END CHANGES FROM PYG #####
+    if x.device.type != "ipu":
+        assert (
+            num_nodes <= max_num_nodes_per_graph
+        ).all(), f"Encountered graphs with {num_nodes.max()} nodes, greater than `max_num_nodes = {max_num_nodes_per_graph}`"
 
     out[idx] = x
     out = out.view([batch_size, max_num_nodes_per_graph] + list(x.size())[1:])
 
-    mask = torch.zeros(batch_size * max_num_nodes_per_graph, dtype=torch.bool, device=x.device)
+    # Create a zero-mask on the right device
+    mask_sz = batch_size * max_num_nodes_per_graph
+    if x.device.type in ("ipu", "xla"):
+        mask = torch.zeros(mask_sz, dtype=torch.bool, device="cpu")
+        mask = mask.to(x.device)
+    else:
+        mask = torch.zeros(mask_sz, dtype=torch.bool, device=x.device)
+
+    ##### END CHANGES FROM PYG #####
+
     mask[idx] = 1
     mask = mask.view(batch_size, max_num_nodes_per_graph)
 
diff --git a/goli/nn/architectures/global_architectures.py b/goli/nn/architectures/global_architectures.py
index d0144ff89..cbfaa95f9 100644
--- a/goli/nn/architectures/global_architectures.py
+++ b/goli/nn/architectures/global_architectures.py
@@ -13,7 +13,7 @@
 
 # goli imports
 from goli.nn.base_layers import FCLayer, get_activation, get_norm
-from goli.nn.base_graph_layer import BaseGraphModule
+from goli.nn.base_graph_layer import BaseGraphModule, BaseGraphStructure
 from goli.nn.residual_connections import (
     ResidualConnectionBase,
     ResidualConnectionWeighted,
@@ -1481,6 +1481,23 @@ def make_mup_base_kwargs(self, divide_factor: float = 2.0) -> Dict[str, Any]:
 
         return kwargs
 
+    def set_max_num_nodes_edges_per_graph(self, max_nodes: Optional[int], max_edges: Optional[int]) -> None:
+        """
+        Set the maximum number of nodes and edges for all gnn layers
+
+        Parameters:
+            max_nodes: Maximum number of nodes in the dataset.
+                This will be useful for certain architecture, but ignored by others.
+
+            max_edges: Maximum number of edges in the dataset.
+                This will be useful for certain architecture, but ignored by others.
+        """
+        if self.gnn is not None:
+            for layer in self.gnn.layers:
+                if isinstance(layer, BaseGraphStructure):
+                    layer.max_num_nodes_per_graph = max_nodes
+                    layer.max_num_edges_per_graph = max_edges
+
     def __repr__(self) -> str:
         r"""
         Controls how the class is printed
@@ -1732,6 +1749,24 @@ def make_mup_base_kwargs(self, divide_factor: float = 2.0) -> Dict[str, Any]:
         )["task_heads_kwargs"]
         return kwargs
 
+    def set_max_num_nodes_edges_per_graph(self, max_nodes: Optional[int], max_edges: Optional[int]) -> None:
+        """
+        Set the maximum number of nodes and edges for all gnn layers
+
+        Parameters:
+            max_nodes: Maximum number of nodes in the dataset.
+                This will be useful for certain architecture, but ignored by others.
+
+            max_edges: Maximum number of edges in the dataset.
+                This will be useful for certain architecture, but ignored by others.
+        """
+        super().set_max_num_nodes_edges_per_graph(max_nodes, max_edges)
+        for task_head in self.task_heads.task_heads.values():
+            for layer in task_head.layers:
+                if isinstance(layer, BaseGraphStructure):
+                    layer.max_num_nodes_per_graph = max_nodes
+                    layer.max_num_edges_per_graph = max_edges
+
     def __repr__(self):
         task_str = self.task_heads.__repr__()
         task_str = "    Task heads:\n    " + "    ".join(task_str.splitlines(True))
diff --git a/goli/nn/base_graph_layer.py b/goli/nn/base_graph_layer.py
index bff3e3190..ef45c47e9 100644
--- a/goli/nn/base_graph_layer.py
+++ b/goli/nn/base_graph_layer.py
@@ -66,9 +66,10 @@ def __init__(
         self.normalization = normalization
         self.dropout = dropout
         self.activation = activation
+        self._max_num_nodes_per_graph = None
+        self._max_num_edges_per_graph = None
 
     def _initialize_activation_dropout_norm(self):
-
         if not isinstance(self, nn.Module):
             raise TypeError(
                 "This function requires the current object to be an `nn.Module`. Use multi-inheritance or the class `BaseGraphModule` instead"
@@ -88,7 +89,6 @@ def _parse_dropout(self, dropout):
         return
 
     def _parse_norm(self, normalization, dim=None):
-
         if dim is None:
             dim = self.out_dim * self.out_dim_factor
         if normalization is None or normalization == "none":
@@ -215,6 +215,42 @@ def out_dim_factor(self) -> int:
         """
         ...
 
+    @property
+    def max_num_nodes_per_graph(self) -> Optional[int]:
+        """
+        Get the maximum number of nodes per graph. Useful for reshaping a compiled model (IPU)
+        """
+        return self._max_num_nodes_per_graph
+
+    @max_num_nodes_per_graph.setter
+    def max_num_nodes_per_graph(self, value: Optional[int]):
+        """
+        Set the maximum number of nodes per graph. Useful for reshaping a compiled model (IPU)
+        """
+        if value is not None:
+            assert isinstance(value, int) and (
+                value > 0
+            ), f"Value should be a positive integer, provided f{value} of type {type(value)}"
+        self._max_num_nodes_per_graph = value
+
+    @property
+    def max_num_edges_per_graph(self) -> Optional[int]:
+        """
+        Get the maximum number of nodes per graph. Useful for reshaping a compiled model (IPU)
+        """
+        return self._max_num_edges_per_graph
+
+    @max_num_edges_per_graph.setter
+    def max_num_edges_per_graph(self, value: Optional[int]):
+        """
+        Set the maximum number of nodes per graph. Useful for reshaping a compiled model (IPU)
+        """
+        if value is not None:
+            assert isinstance(value, int) and (
+                value > 0
+            ), f"Value should be a positive integer, provided f{value} of type {type(value)}"
+        self._max_num_edges_per_graph = value
+
     def __repr__(self):
         r"""
         Controls how the class is printed
@@ -283,8 +319,8 @@ def check_intpus_allow_int(obj, edge_index, size):
     if isinstance(edge_index, Tensor):
         # These 3 lines are different. They check for more int types and avoid overflow
         assert edge_index.dtype in (torch.long, torch.int64, torch.int32, torch.int16)
-        assert edge_index.min() >= 0
-        assert edge_index.max() < torch.iinfo(edge_index.dtype).max
+        # assert edge_index.min() >= 0
+        # assert edge_index.max() < torch.iinfo(edge_index.dtype).max
 
         assert edge_index.dim() == 2
         assert edge_index.size(0) == 2
diff --git a/goli/nn/base_layers.py b/goli/nn/base_layers.py
index fa3ee376e..1e0ee1b76 100644
--- a/goli/nn/base_layers.py
+++ b/goli/nn/base_layers.py
@@ -93,7 +93,6 @@ def _reset_parameters(self):
             mupi.xavier_normal_(self.bias_v)
 
     def forward(self, *args, **kwargs) -> Tuple[Tensor, Optional[Tensor]]:
-
         # Patching the forward to use a different scaling for the dot-product
         prev_fn = F._scaled_dot_product_attention
         F._scaled_dot_product_attention = _mup_scaled_dot_product_attention
@@ -236,7 +235,6 @@ def __init__(
         init_fn: Optional[Callable] = None,
         is_readout_layer: bool = False,
     ):
-
         r"""
         A simple fully connected and customizable layer. This layer is centered around a `torch.nn.Linear` module.
         The order in which transformations are applied is:
diff --git a/goli/nn/encoders/laplace_pos_encoder.py b/goli/nn/encoders/laplace_pos_encoder.py
index de8f48e4d..dd0b0142d 100644
--- a/goli/nn/encoders/laplace_pos_encoder.py
+++ b/goli/nn/encoders/laplace_pos_encoder.py
@@ -90,7 +90,6 @@ def parse_on_keys(self, on_keys):
         return on_keys
 
     def forward(self, eigvals, eigvecs):
-
         # TODO Dom: add random flipping to the Laplacian encoder
         if self.training:
             sign_flip = torch.rand(eigvecs.size(1), device=eigvecs.device)
diff --git a/goli/nn/encoders/signnet_pos_encoder.py b/goli/nn/encoders/signnet_pos_encoder.py
index 0bcb35d28..7fa1ba419 100644
--- a/goli/nn/encoders/signnet_pos_encoder.py
+++ b/goli/nn/encoders/signnet_pos_encoder.py
@@ -281,7 +281,6 @@ def parse_on_keys(self, on_keys):
         return on_keys
 
     def forward(self, eigvecs, edge_index, batch_index):
-
         pos_enc = eigvecs.unsqueeze(-1)  # (Num nodes) x (Num Eigenvectors) x 1
 
         empty_mask = torch.isnan(pos_enc)
diff --git a/goli/nn/pyg_layers/gin_pyg.py b/goli/nn/pyg_layers/gin_pyg.py
index ee785b397..7fca261e8 100644
--- a/goli/nn/pyg_layers/gin_pyg.py
+++ b/goli/nn/pyg_layers/gin_pyg.py
@@ -89,7 +89,6 @@ def __init__(
         self.model.__check_input__ = partial(check_intpus_allow_int, self)
 
     def forward(self, batch: Union[Data, Batch]):
-
         batch.h = self.model(batch.h, batch.edge_index)
         batch.h = self.apply_norm_activation_dropout(batch.h)
 
@@ -229,7 +228,6 @@ def __init__(
         self.model.__check_input__ = partial(check_intpus_allow_int, self)
 
     def forward(self, batch):
-
         batch.h = self.model(batch.h, batch.edge_index, batch.edge_attr)
         batch.h = self.apply_norm_activation_dropout(batch.h)
 
diff --git a/goli/nn/pyg_layers/gps_pyg.py b/goli/nn/pyg_layers/gps_pyg.py
index 2924b5193..f23bbf4fe 100644
--- a/goli/nn/pyg_layers/gps_pyg.py
+++ b/goli/nn/pyg_layers/gps_pyg.py
@@ -11,7 +11,7 @@
 from goli.nn.pyg_layers import GatedGCNPyg, GINConvPyg, GINEConvPyg, PNAMessagePassingPyg
 from goli.utils.decorators import classproperty
 from goli.ipu.to_dense_batch import to_dense_batch, to_sparse_batch
-
+from goli.ipu.ipu_utils import import_poptorch
 
 PYG_LAYERS_DICT = {
     "pyg:gin": GINConvPyg,
@@ -77,7 +77,7 @@ def __init__(
             out_dim=out_dim,
             activation=activation,
             dropout=dropout,
-            normalization=normalization,
+            normalization=None,
         )
 
         # Dropout layers
@@ -94,7 +94,7 @@ def __init__(
         # Normalization layers
         self.norm_layer_local = self._parse_norm(normalization=self.normalization, dim=in_dim)
         self.norm_layer_attn = self._parse_norm(normalization=self.normalization, dim=in_dim)
-        self.norm_layer_ff = self.norm_layer
+        self.norm_layer_ff = self._parse_norm(self.normalization)
 
         # Set the default values for the MPNN layer
         if mpnn_kwargs is None:
@@ -140,17 +140,21 @@ def forward(self, batch):
         # * batch.batch is the indicator vector for nodes of which graph it belongs to
         # * h_dense
         if self.attn_layer is not None:
-
-            # If there's padding, then we are on IPU
-            on_ipu = ("graph_is_true" in batch.keys) and (not batch.graph_is_true.all())
+            # Check whether the model runs on IPU, if so define a maximal number of nodes per graph when reshaping
+            poptorch = import_poptorch(raise_error=False)
+            on_ipu = (poptorch is not None) and (poptorch.isRunningOnIpu())
+            max_num_nodes_per_graph = None
             if on_ipu:
-                max_num_nodes_per_graph = batch.dataset_max_nodes_per_graph[0].item()
-            else:
-                max_num_nodes_per_graph = None
+                max_num_nodes_per_graph = self.max_num_nodes_per_graph
 
             # Convert the tensor to a dense batch, then back to a sparse batch
+            batch_size = None if h.device.type != "ipu" else batch.graph_is_true.shape[0]
             h_dense, mask, idx = to_dense_batch(
-                h, batch.batch, max_num_nodes_per_graph=max_num_nodes_per_graph, drop_nodes_last_graph=on_ipu
+                h,
+                batch=batch.batch,
+                batch_size=batch_size,
+                max_num_nodes_per_graph=max_num_nodes_per_graph,
+                drop_nodes_last_graph=on_ipu,
             )
             h_attn = self._sa_block(h_dense, None, ~mask)
             h_attn = to_sparse_batch(h_attn, idx)
diff --git a/goli/nn/pyg_layers/pna_pyg.py b/goli/nn/pyg_layers/pna_pyg.py
index b42ac4145..a8c8223a9 100644
--- a/goli/nn/pyg_layers/pna_pyg.py
+++ b/goli/nn/pyg_layers/pna_pyg.py
@@ -162,7 +162,6 @@ def forward(self, batch: Union[Data, Batch]):
         return batch
 
     def message(self, x_i: Tensor, x_j: Tensor, edge_attr: OptTensor) -> Tensor:
-
         h: Tensor = x_i  # Dummy.
         if (edge_attr is not None) and (self.edge_encoder is not None):
             edge_attr = self.edge_encoder(edge_attr)
@@ -175,7 +174,6 @@ def message(self, x_i: Tensor, x_j: Tensor, edge_attr: OptTensor) -> Tensor:
     def aggregate(
         self, inputs: Tensor, index: Tensor, edge_index: Tensor, dim_size: Optional[int] = None
     ) -> Tensor:
-
         outs = []
 
         for aggregator in self.aggregators:
diff --git a/goli/nn/residual_connections.py b/goli/nn/residual_connections.py
index b3b8e7757..88f99019d 100644
--- a/goli/nn/residual_connections.py
+++ b/goli/nn/residual_connections.py
@@ -76,11 +76,9 @@ def h_dim_increase_type(cls):
         ...
 
     def get_true_out_dims(self, out_dims: List):
-
         true_out_dims = [out_dims[0]]
         out_dims_at_skip = [out_dims[0]]
         for ii in range(1, len(out_dims) - 1):
-
             # For the `None` type, don't change the output dims
             if self.h_dim_increase_type is None:
                 true_out_dims.append(out_dims[ii])
diff --git a/goli/trainer/metrics.py b/goli/trainer/metrics.py
index 9550bcc56..900483f60 100644
--- a/goli/trainer/metrics.py
+++ b/goli/trainer/metrics.py
@@ -29,7 +29,6 @@ def __init__(
         th_on_target: bool = False,
         target_to_int: bool = False,
     ):
-
         # Basic params
         self.threshold = threshold
         self.th_on_target = th_on_target
diff --git a/goli/trainer/predictor.py b/goli/trainer/predictor.py
index 7770a0f8e..05852c535 100644
--- a/goli/trainer/predictor.py
+++ b/goli/trainer/predictor.py
@@ -14,6 +14,7 @@
 from goli.trainer.predictor_options import EvalOptions, FlagOptions, ModelOptions, OptimOptions
 from goli.trainer.predictor_summaries import TaskSummaries
 from goli.nn.base_graph_layer import get_node_feats, set_node_feats
+from goli.data.datamodule import BaseDataModule
 
 GOLI_PRETRAINED_MODELS = {
     "goli-zinc-micro-dummy-test": "gcs://goli-public/pretrained-models/goli-zinc-micro-dummy-test/model.ckpt"
@@ -177,7 +178,6 @@ def _convert_features_dtype(self, feats):
         return feats
 
     def configure_optimizers(self, impl=None):
-
         if impl is None:
             impl = torch.optim.Adam
 
@@ -308,11 +308,12 @@ def _general_step(self, batch: Dict[str, Tensor], step_name: str, to_cpu: bool)
         for task in self.tasks:
             step_dict[
                 self.task_epoch_summary.metric_log_name(task, self.loss_fun[task]._get_name(), step_name)
-            ] = loss.detach().cpu()
+            ] = loss.detach()
 
         step_dict["loss"] = loss
         # print("loss ", self.global_step, self.current_epoch, loss)
         step_dict["task_losses"] = task_losses
+        step_dict["gradient_norm"] = self.get_gradient_norm()
         return step_dict
 
     def flag_step(self, batch: Dict[str, Tensor], step_name: str, to_cpu: bool) -> Dict[str, Any]:
@@ -384,10 +385,28 @@ def flag_step(self, batch: Dict[str, Tensor], step_name: str, to_cpu: bool) -> D
         step_dict["task_losses"] = task_losses
         return step_dict
 
-    def on_train_batch_end(self, outputs, batch: Any, batch_idx: int, unused: int = 0) -> None:
+    def on_train_batch_end(self, outputs, batch: Any, batch_idx: int) -> None:
+        # this code is likely repeated for validation and testing, this should be moved to a function
+        self.task_epoch_summary.update_predictor_state(
+            step_name="train",
+            targets=outputs["targets"],
+            predictions=outputs["preds"],
+            loss=outputs["loss"],  # This is the weighted loss for now, but change to task-specific loss
+            task_losses=outputs["task_losses"],
+            n_epochs=self.current_epoch,
+        )
+        metrics_logs = self.task_epoch_summary.get_metrics_logs()  # Dict[task, metric_logs]
+        metrics_logs["_global"]["grad_norm"] = self.get_gradient_norm()
+        outputs.update(metrics_logs)  # Dict[task, metric_logs]. Concatenate them?
+
+        concatenated_metrics_logs = self.task_epoch_summary.concatenate_metrics_logs(metrics_logs)
+        concatenated_metrics_logs["loss"] = outputs["loss"]
+        outputs["grad_norm"] = self.get_gradient_norm()
+        concatenated_metrics_logs["train/grad_norm"] = outputs["grad_norm"]
+
         if self.logger is not None:
             self.logger.log_metrics(
-                outputs, step=self.global_step
+                concatenated_metrics_logs, step=self.global_step
             )  # This is a pytorch lightning function call
 
     def training_step(self, batch: Dict[str, Tensor], to_cpu: bool = True) -> Dict[str, Any]:
@@ -398,37 +417,10 @@ def training_step(self, batch: Dict[str, Tensor], to_cpu: bool = True) -> Dict[s
             step_dict = self.flag_step(batch=batch, step_name="train", to_cpu=to_cpu)
         # Train normally, without using FLAG
         elif self.flag_kwargs["n_steps"] == 0:
-            step_dict = self._general_step(batch=batch, step_name="train", to_cpu=True)
+            # step_dict = self._general_step(batch=batch, step_name="train", to_cpu=True)
+            step_dict = self._general_step(batch=batch, step_name="train", to_cpu=to_cpu)
 
-        self.task_epoch_summary.update_predictor_state(
-            step_name="train",
-            targets=step_dict["targets"],
-            predictions=step_dict["preds"],
-            loss=step_dict["loss"],  # This is the weighted loss for now, but change to task-sepcific loss
-            task_losses=step_dict["task_losses"],
-            n_epochs=self.current_epoch,
-        )
-        metrics_logs = self.task_epoch_summary.get_metrics_logs()  # Dict[task, metric_logs]
-        metrics_logs["_global"]["grad_norm"] = self.get_gradient_norm()
-        step_dict.update(metrics_logs)  # Dict[task, metric_logs]. Concatenate them?
-
-        concatenated_metrics_logs = self.task_epoch_summary.concatenate_metrics_logs(metrics_logs)
-        concatenated_metrics_logs["loss"] = step_dict["loss"]
-        step_dict["grad_norm"] = self.get_gradient_norm()
-        concatenated_metrics_logs["train/grad_norm"] = step_dict["grad_norm"]
-
-        # Predictions and targets are no longer needed after the step.
-        # Keeping them will increase memory usage significantly for large datasets.
-        step_dict.pop("preds")
-        step_dict.pop("targets")
-        step_dict.pop("weights")
-
-        return concatenated_metrics_logs  # Returning the metrics_logs with the loss
-
-    def on_train_batch_end(self, outputs, batch: Any, batch_idx: int, unused: int = 0) -> None:
-        # Wandb metric tracking here
-        if self.logger is not None:
-            self.logger.log_metrics(outputs, step=self.global_step)
+        return step_dict  # Returning the metrics_logs with the loss
 
     def get_gradient_norm(self):
         # compute the norm
@@ -451,9 +443,11 @@ def _general_epoch_end(self, outputs: Dict[str, Any], step_name: str) -> None:
         # Transform the list of dict of dict, into a dict of list of dict
         preds = {}
         targets = {}
+        device = device = outputs[0]["preds"][self.tasks[0]].device  # should be better way to do this
+        # device = 0
         for task in self.tasks:
-            preds[task] = torch.cat([out["preds"][task] for out in outputs], dim=0)
-            targets[task] = torch.cat([out["targets"][task] for out in outputs], dim=0)
+            preds[task] = torch.cat([out["preds"][task].to(device=device) for out in outputs], dim=0)
+            targets[task] = torch.cat([out["targets"][task].to(device=device) for out in outputs], dim=0)
         if ("weights" in outputs[0].keys()) and (outputs[0]["weights"] is not None):
             weights = torch.cat([out["weights"] for out in outputs], dim=0)
         else:
@@ -489,7 +483,6 @@ def training_epoch_end(self, outputs: Dict):
         pass
 
     def validation_epoch_end(self, outputs: Dict[str, Any]):
-
         metrics_logs = self._general_epoch_end(outputs=outputs, step_name="val")
         concatenated_metrics_logs = self.task_epoch_summary.concatenate_metrics_logs(metrics_logs)
 
@@ -503,7 +496,6 @@ def validation_epoch_end(self, outputs: Dict[str, Any]):
         full_dict.update(self.task_epoch_summary.get_dict_summary())
 
     def test_epoch_end(self, outputs: Dict[str, Any]):
-
         metrics_logs = self._general_epoch_end(outputs=outputs, step_name="test")
         concatenated_metrics_logs = self.task_epoch_summary.concatenate_metrics_logs(metrics_logs)
 
@@ -559,3 +551,11 @@ def load_pretrained_models(name: str):
             )
 
         return PredictorModule.load_from_checkpoint(GOLI_PRETRAINED_MODELS[name])
+
+    def set_max_nodes_edges_per_graph(self, datamodule: BaseDataModule, stages: Optional[List[str]] = None):
+        datamodule.setup()
+
+        max_nodes = datamodule.get_max_num_nodes_datamodule(stages)
+        max_edges = datamodule.get_max_num_edges_datamodule(stages)
+
+        self.model.set_max_num_nodes_edges_per_graph(max_nodes, max_edges)
diff --git a/goli/trainer/predictor_summaries.py b/goli/trainer/predictor_summaries.py
index fd1b0e4fc..ece45e988 100644
--- a/goli/trainer/predictor_summaries.py
+++ b/goli/trainer/predictor_summaries.py
@@ -30,6 +30,7 @@ def get_metrics_logs(self, **kwargs):
 
 class Summary(SummaryInterface):
     r"""A container to be used by the Predictor Module that stores the results for the given metrics on the predictions and targets provided."""
+
     # TODO (Gabriela): Default argument cannot be []
     def __init__(
         self,
diff --git a/goli/utils/dict_tensor.py b/goli/utils/dict_tensor.py
index 27b1ababe..0fd8f1e70 100644
--- a/goli/utils/dict_tensor.py
+++ b/goli/utils/dict_tensor.py
@@ -118,7 +118,6 @@ def _dict_func_wrapper(self, func: Callable):
         """
 
         def wrap(first_input: Union[DictTensor, Any], *args, **kwargs):
-
             # The following lines are for when a function is 'bounded' and
             # is called form `torch.FUNCTION()` instead of `Tensor.FUNCTION()`.
             this_func, this_args, this_kwargs = func, args, kwargs
diff --git a/goli/utils/fs.py b/goli/utils/fs.py
index 26ab3b2a8..2ae1c1102 100644
--- a/goli/utils/fs.py
+++ b/goli/utils/fs.py
@@ -159,7 +159,6 @@ def copy(
         source_file = source
 
     if isinstance(destination, (str, os.PathLike)):
-
         # adapt the file mode of the destination depending on the source file.
         destination_mode = "wb"
         if hasattr(source_file, "mode"):
@@ -181,7 +180,6 @@ def copy(
 
     with source_file as source_stream:
         with destination_file as destination_stream:
-
             if chunk_size is None:
                 # copy without chunks
                 destination_stream.write(source_stream.read())
diff --git a/goli/utils/tensor.py b/goli/utils/tensor.py
index 0bfc44de9..25d2d3b58 100644
--- a/goli/utils/tensor.py
+++ b/goli/utils/tensor.py
@@ -14,7 +14,6 @@
 
 
 def save_im(im_dir, im_name: str, ext: List[str] = ["svg", "png"], dpi: int = 600) -> None:
-
     if not os.path.exists(im_dir):
         if im_dir[-1] not in ["/", "\\"]:
             im_dir += "/"
diff --git a/mup_base_params.yaml b/mup_base_params.yaml
index b5e0d1ce3..20a085308 100644
--- a/mup_base_params.yaml
+++ b/mup_base_params.yaml
@@ -22,9 +22,9 @@ gnn.layers.0.ff_linear2.linear.weight:
 - 16
 - 32
 gnn.layers.0.ff_out.linear.bias:
-- 32
+- 250
 gnn.layers.0.ff_out.linear.weight:
-- 32
+- 250
 - 16
 gnn.layers.0.mpnn.model.lin.bias:
 - 16
@@ -41,807 +41,615 @@ gnn.layers.0.mpnn.model.nn.fully_connected.1.linear.bias:
 gnn.layers.0.mpnn.model.nn.fully_connected.1.linear.weight:
 - 16
 - 16
-gnn.layers.0.norm_layer.bias:
-- 32
-gnn.layers.0.norm_layer.weight:
-- 32
-gnn.layers.0.norm_layer_attn.bias:
-- 16
-gnn.layers.0.norm_layer_attn.weight:
-- 16
-gnn.layers.0.norm_layer_local.bias:
-- 16
-gnn.layers.0.norm_layer_local.weight:
-- 16
 gnn.layers.1.attn_layer.in_proj_bias:
-- 96
+- 750
 gnn.layers.1.attn_layer.in_proj_weight:
-- 96
-- 32
+- 750
+- 250
 gnn.layers.1.attn_layer.out_proj.bias:
-- 32
+- 250
 gnn.layers.1.attn_layer.out_proj.weight:
-- 32
-- 32
+- 250
+- 250
 gnn.layers.1.ff_linear1.linear.bias:
-- 64
+- 500
 gnn.layers.1.ff_linear1.linear.weight:
-- 64
-- 32
+- 500
+- 250
 gnn.layers.1.ff_linear2.linear.bias:
-- 32
+- 250
 gnn.layers.1.ff_linear2.linear.weight:
-- 32
-- 64
+- 250
+- 500
 gnn.layers.1.ff_out.linear.bias:
-- 32
+- 250
 gnn.layers.1.ff_out.linear.weight:
-- 32
-- 32
+- 250
+- 250
 gnn.layers.1.mpnn.model.lin.bias:
-- 32
+- 250
 gnn.layers.1.mpnn.model.lin.weight:
-- 32
+- 250
 - 8
 gnn.layers.1.mpnn.model.nn.fully_connected.0.linear.bias:
-- 32
+- 250
 gnn.layers.1.mpnn.model.nn.fully_connected.0.linear.weight:
-- 32
-- 32
+- 250
+- 250
 gnn.layers.1.mpnn.model.nn.fully_connected.1.linear.bias:
-- 32
+- 250
 gnn.layers.1.mpnn.model.nn.fully_connected.1.linear.weight:
-- 32
-- 32
-gnn.layers.1.norm_layer.bias:
-- 32
-gnn.layers.1.norm_layer.weight:
-- 32
-gnn.layers.1.norm_layer_attn.bias:
-- 32
-gnn.layers.1.norm_layer_attn.weight:
-- 32
-gnn.layers.1.norm_layer_local.bias:
-- 32
-gnn.layers.1.norm_layer_local.weight:
-- 32
+- 250
+- 250
 gnn.layers.10.attn_layer.in_proj_bias:
-- 96
+- 750
 gnn.layers.10.attn_layer.in_proj_weight:
-- 96
-- 32
+- 750
+- 250
 gnn.layers.10.attn_layer.out_proj.bias:
-- 32
+- 250
 gnn.layers.10.attn_layer.out_proj.weight:
-- 32
-- 32
+- 250
+- 250
 gnn.layers.10.ff_linear1.linear.bias:
-- 64
+- 500
 gnn.layers.10.ff_linear1.linear.weight:
-- 64
-- 32
+- 500
+- 250
 gnn.layers.10.ff_linear2.linear.bias:
-- 32
+- 250
 gnn.layers.10.ff_linear2.linear.weight:
-- 32
-- 64
+- 250
+- 500
 gnn.layers.10.ff_out.linear.bias:
-- 32
+- 250
 gnn.layers.10.ff_out.linear.weight:
-- 32
-- 32
+- 250
+- 250
 gnn.layers.10.mpnn.model.lin.bias:
-- 32
+- 250
 gnn.layers.10.mpnn.model.lin.weight:
-- 32
+- 250
 - 8
 gnn.layers.10.mpnn.model.nn.fully_connected.0.linear.bias:
-- 32
+- 250
 gnn.layers.10.mpnn.model.nn.fully_connected.0.linear.weight:
-- 32
-- 32
+- 250
+- 250
 gnn.layers.10.mpnn.model.nn.fully_connected.1.linear.bias:
-- 32
+- 250
 gnn.layers.10.mpnn.model.nn.fully_connected.1.linear.weight:
-- 32
-- 32
-gnn.layers.10.norm_layer.bias:
-- 32
-gnn.layers.10.norm_layer.weight:
-- 32
-gnn.layers.10.norm_layer_attn.bias:
-- 32
-gnn.layers.10.norm_layer_attn.weight:
-- 32
-gnn.layers.10.norm_layer_local.bias:
-- 32
-gnn.layers.10.norm_layer_local.weight:
-- 32
+- 250
+- 250
 gnn.layers.11.attn_layer.in_proj_bias:
-- 96
+- 750
 gnn.layers.11.attn_layer.in_proj_weight:
-- 96
-- 32
+- 750
+- 250
 gnn.layers.11.attn_layer.out_proj.bias:
-- 32
+- 250
 gnn.layers.11.attn_layer.out_proj.weight:
-- 32
-- 32
+- 250
+- 250
 gnn.layers.11.ff_linear1.linear.bias:
-- 64
+- 500
 gnn.layers.11.ff_linear1.linear.weight:
-- 64
-- 32
+- 500
+- 250
 gnn.layers.11.ff_linear2.linear.bias:
-- 32
+- 250
 gnn.layers.11.ff_linear2.linear.weight:
-- 32
-- 64
+- 250
+- 500
 gnn.layers.11.ff_out.linear.bias:
-- 32
+- 250
 gnn.layers.11.ff_out.linear.weight:
-- 32
-- 32
+- 250
+- 250
 gnn.layers.11.mpnn.model.lin.bias:
-- 32
+- 250
 gnn.layers.11.mpnn.model.lin.weight:
-- 32
+- 250
 - 8
 gnn.layers.11.mpnn.model.nn.fully_connected.0.linear.bias:
-- 32
+- 250
 gnn.layers.11.mpnn.model.nn.fully_connected.0.linear.weight:
-- 32
-- 32
+- 250
+- 250
 gnn.layers.11.mpnn.model.nn.fully_connected.1.linear.bias:
-- 32
+- 250
 gnn.layers.11.mpnn.model.nn.fully_connected.1.linear.weight:
-- 32
-- 32
-gnn.layers.11.norm_layer.bias:
-- 32
-gnn.layers.11.norm_layer.weight:
-- 32
-gnn.layers.11.norm_layer_attn.bias:
-- 32
-gnn.layers.11.norm_layer_attn.weight:
-- 32
-gnn.layers.11.norm_layer_local.bias:
-- 32
-gnn.layers.11.norm_layer_local.weight:
-- 32
+- 250
+- 250
 gnn.layers.12.attn_layer.in_proj_bias:
-- 96
+- 750
 gnn.layers.12.attn_layer.in_proj_weight:
-- 96
-- 32
+- 750
+- 250
 gnn.layers.12.attn_layer.out_proj.bias:
-- 32
+- 250
 gnn.layers.12.attn_layer.out_proj.weight:
-- 32
-- 32
+- 250
+- 250
 gnn.layers.12.ff_linear1.linear.bias:
-- 64
+- 500
 gnn.layers.12.ff_linear1.linear.weight:
-- 64
-- 32
+- 500
+- 250
 gnn.layers.12.ff_linear2.linear.bias:
-- 32
+- 250
 gnn.layers.12.ff_linear2.linear.weight:
-- 32
-- 64
+- 250
+- 500
 gnn.layers.12.ff_out.linear.bias:
-- 32
+- 250
 gnn.layers.12.ff_out.linear.weight:
-- 32
-- 32
+- 250
+- 250
 gnn.layers.12.mpnn.model.lin.bias:
-- 32
+- 250
 gnn.layers.12.mpnn.model.lin.weight:
-- 32
+- 250
 - 8
 gnn.layers.12.mpnn.model.nn.fully_connected.0.linear.bias:
-- 32
+- 250
 gnn.layers.12.mpnn.model.nn.fully_connected.0.linear.weight:
-- 32
-- 32
+- 250
+- 250
 gnn.layers.12.mpnn.model.nn.fully_connected.1.linear.bias:
-- 32
+- 250
 gnn.layers.12.mpnn.model.nn.fully_connected.1.linear.weight:
-- 32
-- 32
-gnn.layers.12.norm_layer.bias:
-- 32
-gnn.layers.12.norm_layer.weight:
-- 32
-gnn.layers.12.norm_layer_attn.bias:
-- 32
-gnn.layers.12.norm_layer_attn.weight:
-- 32
-gnn.layers.12.norm_layer_local.bias:
-- 32
-gnn.layers.12.norm_layer_local.weight:
-- 32
+- 250
+- 250
 gnn.layers.13.attn_layer.in_proj_bias:
-- 96
+- 750
 gnn.layers.13.attn_layer.in_proj_weight:
-- 96
-- 32
+- 750
+- 250
 gnn.layers.13.attn_layer.out_proj.bias:
-- 32
+- 250
 gnn.layers.13.attn_layer.out_proj.weight:
-- 32
-- 32
+- 250
+- 250
 gnn.layers.13.ff_linear1.linear.bias:
-- 64
+- 500
 gnn.layers.13.ff_linear1.linear.weight:
-- 64
-- 32
+- 500
+- 250
 gnn.layers.13.ff_linear2.linear.bias:
-- 32
+- 250
 gnn.layers.13.ff_linear2.linear.weight:
-- 32
-- 64
+- 250
+- 500
 gnn.layers.13.ff_out.linear.bias:
-- 32
+- 250
 gnn.layers.13.ff_out.linear.weight:
-- 32
-- 32
+- 250
+- 250
 gnn.layers.13.mpnn.model.lin.bias:
-- 32
+- 250
 gnn.layers.13.mpnn.model.lin.weight:
-- 32
+- 250
 - 8
 gnn.layers.13.mpnn.model.nn.fully_connected.0.linear.bias:
-- 32
+- 250
 gnn.layers.13.mpnn.model.nn.fully_connected.0.linear.weight:
-- 32
-- 32
+- 250
+- 250
 gnn.layers.13.mpnn.model.nn.fully_connected.1.linear.bias:
-- 32
+- 250
 gnn.layers.13.mpnn.model.nn.fully_connected.1.linear.weight:
-- 32
-- 32
-gnn.layers.13.norm_layer.bias:
-- 32
-gnn.layers.13.norm_layer.weight:
-- 32
-gnn.layers.13.norm_layer_attn.bias:
-- 32
-gnn.layers.13.norm_layer_attn.weight:
-- 32
-gnn.layers.13.norm_layer_local.bias:
-- 32
-gnn.layers.13.norm_layer_local.weight:
-- 32
+- 250
+- 250
 gnn.layers.14.attn_layer.in_proj_bias:
-- 96
+- 750
 gnn.layers.14.attn_layer.in_proj_weight:
-- 96
-- 32
+- 750
+- 250
 gnn.layers.14.attn_layer.out_proj.bias:
-- 32
+- 250
 gnn.layers.14.attn_layer.out_proj.weight:
-- 32
-- 32
+- 250
+- 250
 gnn.layers.14.ff_linear1.linear.bias:
-- 64
+- 500
 gnn.layers.14.ff_linear1.linear.weight:
-- 64
-- 32
+- 500
+- 250
 gnn.layers.14.ff_linear2.linear.bias:
-- 32
+- 250
 gnn.layers.14.ff_linear2.linear.weight:
-- 32
-- 64
+- 250
+- 500
 gnn.layers.14.ff_out.linear.bias:
-- 32
+- 250
 gnn.layers.14.ff_out.linear.weight:
-- 32
-- 32
+- 250
+- 250
 gnn.layers.14.mpnn.model.lin.bias:
-- 32
+- 250
 gnn.layers.14.mpnn.model.lin.weight:
-- 32
+- 250
 - 8
 gnn.layers.14.mpnn.model.nn.fully_connected.0.linear.bias:
-- 32
+- 250
 gnn.layers.14.mpnn.model.nn.fully_connected.0.linear.weight:
-- 32
-- 32
+- 250
+- 250
 gnn.layers.14.mpnn.model.nn.fully_connected.1.linear.bias:
-- 32
+- 250
 gnn.layers.14.mpnn.model.nn.fully_connected.1.linear.weight:
-- 32
-- 32
-gnn.layers.14.norm_layer.bias:
-- 32
-gnn.layers.14.norm_layer.weight:
-- 32
-gnn.layers.14.norm_layer_attn.bias:
-- 32
-gnn.layers.14.norm_layer_attn.weight:
-- 32
-gnn.layers.14.norm_layer_local.bias:
-- 32
-gnn.layers.14.norm_layer_local.weight:
-- 32
+- 250
+- 250
 gnn.layers.15.attn_layer.in_proj_bias:
-- 96
+- 750
 gnn.layers.15.attn_layer.in_proj_weight:
-- 96
-- 32
+- 750
+- 250
 gnn.layers.15.attn_layer.out_proj.bias:
-- 32
+- 250
 gnn.layers.15.attn_layer.out_proj.weight:
-- 32
-- 32
+- 250
+- 250
 gnn.layers.15.ff_linear1.linear.bias:
-- 64
+- 500
 gnn.layers.15.ff_linear1.linear.weight:
-- 64
-- 32
+- 500
+- 250
 gnn.layers.15.ff_linear2.linear.bias:
-- 32
+- 250
 gnn.layers.15.ff_linear2.linear.weight:
-- 32
-- 64
+- 250
+- 500
 gnn.layers.15.ff_out.linear.bias:
-- 16
+- 250
 gnn.layers.15.ff_out.linear.weight:
-- 16
-- 32
+- 250
+- 250
 gnn.layers.15.mpnn.model.lin.bias:
-- 32
+- 250
 gnn.layers.15.mpnn.model.lin.weight:
-- 32
+- 250
 - 8
 gnn.layers.15.mpnn.model.nn.fully_connected.0.linear.bias:
-- 32
+- 250
 gnn.layers.15.mpnn.model.nn.fully_connected.0.linear.weight:
-- 32
-- 32
+- 250
+- 250
 gnn.layers.15.mpnn.model.nn.fully_connected.1.linear.bias:
-- 32
+- 250
 gnn.layers.15.mpnn.model.nn.fully_connected.1.linear.weight:
-- 32
-- 32
-gnn.layers.15.norm_layer.bias:
-- 16
-gnn.layers.15.norm_layer.weight:
-- 16
-gnn.layers.15.norm_layer_attn.bias:
-- 32
-gnn.layers.15.norm_layer_attn.weight:
-- 32
-gnn.layers.15.norm_layer_local.bias:
-- 32
-gnn.layers.15.norm_layer_local.weight:
-- 32
+- 250
+- 250
 gnn.layers.2.attn_layer.in_proj_bias:
-- 96
+- 750
 gnn.layers.2.attn_layer.in_proj_weight:
-- 96
-- 32
+- 750
+- 250
 gnn.layers.2.attn_layer.out_proj.bias:
-- 32
+- 250
 gnn.layers.2.attn_layer.out_proj.weight:
-- 32
-- 32
+- 250
+- 250
 gnn.layers.2.ff_linear1.linear.bias:
-- 64
+- 500
 gnn.layers.2.ff_linear1.linear.weight:
-- 64
-- 32
+- 500
+- 250
 gnn.layers.2.ff_linear2.linear.bias:
-- 32
+- 250
 gnn.layers.2.ff_linear2.linear.weight:
-- 32
-- 64
+- 250
+- 500
 gnn.layers.2.ff_out.linear.bias:
-- 32
+- 250
 gnn.layers.2.ff_out.linear.weight:
-- 32
-- 32
+- 250
+- 250
 gnn.layers.2.mpnn.model.lin.bias:
-- 32
+- 250
 gnn.layers.2.mpnn.model.lin.weight:
-- 32
+- 250
 - 8
 gnn.layers.2.mpnn.model.nn.fully_connected.0.linear.bias:
-- 32
+- 250
 gnn.layers.2.mpnn.model.nn.fully_connected.0.linear.weight:
-- 32
-- 32
+- 250
+- 250
 gnn.layers.2.mpnn.model.nn.fully_connected.1.linear.bias:
-- 32
+- 250
 gnn.layers.2.mpnn.model.nn.fully_connected.1.linear.weight:
-- 32
-- 32
-gnn.layers.2.norm_layer.bias:
-- 32
-gnn.layers.2.norm_layer.weight:
-- 32
-gnn.layers.2.norm_layer_attn.bias:
-- 32
-gnn.layers.2.norm_layer_attn.weight:
-- 32
-gnn.layers.2.norm_layer_local.bias:
-- 32
-gnn.layers.2.norm_layer_local.weight:
-- 32
+- 250
+- 250
 gnn.layers.3.attn_layer.in_proj_bias:
-- 96
+- 750
 gnn.layers.3.attn_layer.in_proj_weight:
-- 96
-- 32
+- 750
+- 250
 gnn.layers.3.attn_layer.out_proj.bias:
-- 32
+- 250
 gnn.layers.3.attn_layer.out_proj.weight:
-- 32
-- 32
+- 250
+- 250
 gnn.layers.3.ff_linear1.linear.bias:
-- 64
+- 500
 gnn.layers.3.ff_linear1.linear.weight:
-- 64
-- 32
+- 500
+- 250
 gnn.layers.3.ff_linear2.linear.bias:
-- 32
+- 250
 gnn.layers.3.ff_linear2.linear.weight:
-- 32
-- 64
+- 250
+- 500
 gnn.layers.3.ff_out.linear.bias:
-- 32
+- 250
 gnn.layers.3.ff_out.linear.weight:
-- 32
-- 32
+- 250
+- 250
 gnn.layers.3.mpnn.model.lin.bias:
-- 32
+- 250
 gnn.layers.3.mpnn.model.lin.weight:
-- 32
+- 250
 - 8
 gnn.layers.3.mpnn.model.nn.fully_connected.0.linear.bias:
-- 32
+- 250
 gnn.layers.3.mpnn.model.nn.fully_connected.0.linear.weight:
-- 32
-- 32
+- 250
+- 250
 gnn.layers.3.mpnn.model.nn.fully_connected.1.linear.bias:
-- 32
+- 250
 gnn.layers.3.mpnn.model.nn.fully_connected.1.linear.weight:
-- 32
-- 32
-gnn.layers.3.norm_layer.bias:
-- 32
-gnn.layers.3.norm_layer.weight:
-- 32
-gnn.layers.3.norm_layer_attn.bias:
-- 32
-gnn.layers.3.norm_layer_attn.weight:
-- 32
-gnn.layers.3.norm_layer_local.bias:
-- 32
-gnn.layers.3.norm_layer_local.weight:
-- 32
+- 250
+- 250
 gnn.layers.4.attn_layer.in_proj_bias:
-- 96
+- 750
 gnn.layers.4.attn_layer.in_proj_weight:
-- 96
-- 32
+- 750
+- 250
 gnn.layers.4.attn_layer.out_proj.bias:
-- 32
+- 250
 gnn.layers.4.attn_layer.out_proj.weight:
-- 32
-- 32
+- 250
+- 250
 gnn.layers.4.ff_linear1.linear.bias:
-- 64
+- 500
 gnn.layers.4.ff_linear1.linear.weight:
-- 64
-- 32
+- 500
+- 250
 gnn.layers.4.ff_linear2.linear.bias:
-- 32
+- 250
 gnn.layers.4.ff_linear2.linear.weight:
-- 32
-- 64
+- 250
+- 500
 gnn.layers.4.ff_out.linear.bias:
-- 32
+- 250
 gnn.layers.4.ff_out.linear.weight:
-- 32
-- 32
+- 250
+- 250
 gnn.layers.4.mpnn.model.lin.bias:
-- 32
+- 250
 gnn.layers.4.mpnn.model.lin.weight:
-- 32
+- 250
 - 8
 gnn.layers.4.mpnn.model.nn.fully_connected.0.linear.bias:
-- 32
+- 250
 gnn.layers.4.mpnn.model.nn.fully_connected.0.linear.weight:
-- 32
-- 32
+- 250
+- 250
 gnn.layers.4.mpnn.model.nn.fully_connected.1.linear.bias:
-- 32
+- 250
 gnn.layers.4.mpnn.model.nn.fully_connected.1.linear.weight:
-- 32
-- 32
-gnn.layers.4.norm_layer.bias:
-- 32
-gnn.layers.4.norm_layer.weight:
-- 32
-gnn.layers.4.norm_layer_attn.bias:
-- 32
-gnn.layers.4.norm_layer_attn.weight:
-- 32
-gnn.layers.4.norm_layer_local.bias:
-- 32
-gnn.layers.4.norm_layer_local.weight:
-- 32
+- 250
+- 250
 gnn.layers.5.attn_layer.in_proj_bias:
-- 96
+- 750
 gnn.layers.5.attn_layer.in_proj_weight:
-- 96
-- 32
+- 750
+- 250
 gnn.layers.5.attn_layer.out_proj.bias:
-- 32
+- 250
 gnn.layers.5.attn_layer.out_proj.weight:
-- 32
-- 32
+- 250
+- 250
 gnn.layers.5.ff_linear1.linear.bias:
-- 64
+- 500
 gnn.layers.5.ff_linear1.linear.weight:
-- 64
-- 32
+- 500
+- 250
 gnn.layers.5.ff_linear2.linear.bias:
-- 32
+- 250
 gnn.layers.5.ff_linear2.linear.weight:
-- 32
-- 64
+- 250
+- 500
 gnn.layers.5.ff_out.linear.bias:
-- 32
+- 250
 gnn.layers.5.ff_out.linear.weight:
-- 32
-- 32
+- 250
+- 250
 gnn.layers.5.mpnn.model.lin.bias:
-- 32
+- 250
 gnn.layers.5.mpnn.model.lin.weight:
-- 32
+- 250
 - 8
 gnn.layers.5.mpnn.model.nn.fully_connected.0.linear.bias:
-- 32
+- 250
 gnn.layers.5.mpnn.model.nn.fully_connected.0.linear.weight:
-- 32
-- 32
+- 250
+- 250
 gnn.layers.5.mpnn.model.nn.fully_connected.1.linear.bias:
-- 32
+- 250
 gnn.layers.5.mpnn.model.nn.fully_connected.1.linear.weight:
-- 32
-- 32
-gnn.layers.5.norm_layer.bias:
-- 32
-gnn.layers.5.norm_layer.weight:
-- 32
-gnn.layers.5.norm_layer_attn.bias:
-- 32
-gnn.layers.5.norm_layer_attn.weight:
-- 32
-gnn.layers.5.norm_layer_local.bias:
-- 32
-gnn.layers.5.norm_layer_local.weight:
-- 32
+- 250
+- 250
 gnn.layers.6.attn_layer.in_proj_bias:
-- 96
+- 750
 gnn.layers.6.attn_layer.in_proj_weight:
-- 96
-- 32
+- 750
+- 250
 gnn.layers.6.attn_layer.out_proj.bias:
-- 32
+- 250
 gnn.layers.6.attn_layer.out_proj.weight:
-- 32
-- 32
+- 250
+- 250
 gnn.layers.6.ff_linear1.linear.bias:
-- 64
+- 500
 gnn.layers.6.ff_linear1.linear.weight:
-- 64
-- 32
+- 500
+- 250
 gnn.layers.6.ff_linear2.linear.bias:
-- 32
+- 250
 gnn.layers.6.ff_linear2.linear.weight:
-- 32
-- 64
+- 250
+- 500
 gnn.layers.6.ff_out.linear.bias:
-- 32
+- 250
 gnn.layers.6.ff_out.linear.weight:
-- 32
-- 32
+- 250
+- 250
 gnn.layers.6.mpnn.model.lin.bias:
-- 32
+- 250
 gnn.layers.6.mpnn.model.lin.weight:
-- 32
+- 250
 - 8
 gnn.layers.6.mpnn.model.nn.fully_connected.0.linear.bias:
-- 32
+- 250
 gnn.layers.6.mpnn.model.nn.fully_connected.0.linear.weight:
-- 32
-- 32
+- 250
+- 250
 gnn.layers.6.mpnn.model.nn.fully_connected.1.linear.bias:
-- 32
+- 250
 gnn.layers.6.mpnn.model.nn.fully_connected.1.linear.weight:
-- 32
-- 32
-gnn.layers.6.norm_layer.bias:
-- 32
-gnn.layers.6.norm_layer.weight:
-- 32
-gnn.layers.6.norm_layer_attn.bias:
-- 32
-gnn.layers.6.norm_layer_attn.weight:
-- 32
-gnn.layers.6.norm_layer_local.bias:
-- 32
-gnn.layers.6.norm_layer_local.weight:
-- 32
+- 250
+- 250
 gnn.layers.7.attn_layer.in_proj_bias:
-- 96
+- 750
 gnn.layers.7.attn_layer.in_proj_weight:
-- 96
-- 32
+- 750
+- 250
 gnn.layers.7.attn_layer.out_proj.bias:
-- 32
+- 250
 gnn.layers.7.attn_layer.out_proj.weight:
-- 32
-- 32
+- 250
+- 250
 gnn.layers.7.ff_linear1.linear.bias:
-- 64
+- 500
 gnn.layers.7.ff_linear1.linear.weight:
-- 64
-- 32
+- 500
+- 250
 gnn.layers.7.ff_linear2.linear.bias:
-- 32
+- 250
 gnn.layers.7.ff_linear2.linear.weight:
-- 32
-- 64
+- 250
+- 500
 gnn.layers.7.ff_out.linear.bias:
-- 32
+- 250
 gnn.layers.7.ff_out.linear.weight:
-- 32
-- 32
+- 250
+- 250
 gnn.layers.7.mpnn.model.lin.bias:
-- 32
+- 250
 gnn.layers.7.mpnn.model.lin.weight:
-- 32
+- 250
 - 8
 gnn.layers.7.mpnn.model.nn.fully_connected.0.linear.bias:
-- 32
+- 250
 gnn.layers.7.mpnn.model.nn.fully_connected.0.linear.weight:
-- 32
-- 32
+- 250
+- 250
 gnn.layers.7.mpnn.model.nn.fully_connected.1.linear.bias:
-- 32
+- 250
 gnn.layers.7.mpnn.model.nn.fully_connected.1.linear.weight:
-- 32
-- 32
-gnn.layers.7.norm_layer.bias:
-- 32
-gnn.layers.7.norm_layer.weight:
-- 32
-gnn.layers.7.norm_layer_attn.bias:
-- 32
-gnn.layers.7.norm_layer_attn.weight:
-- 32
-gnn.layers.7.norm_layer_local.bias:
-- 32
-gnn.layers.7.norm_layer_local.weight:
-- 32
+- 250
+- 250
 gnn.layers.8.attn_layer.in_proj_bias:
-- 96
+- 750
 gnn.layers.8.attn_layer.in_proj_weight:
-- 96
-- 32
+- 750
+- 250
 gnn.layers.8.attn_layer.out_proj.bias:
-- 32
+- 250
 gnn.layers.8.attn_layer.out_proj.weight:
-- 32
-- 32
+- 250
+- 250
 gnn.layers.8.ff_linear1.linear.bias:
-- 64
+- 500
 gnn.layers.8.ff_linear1.linear.weight:
-- 64
-- 32
+- 500
+- 250
 gnn.layers.8.ff_linear2.linear.bias:
-- 32
+- 250
 gnn.layers.8.ff_linear2.linear.weight:
-- 32
-- 64
+- 250
+- 500
 gnn.layers.8.ff_out.linear.bias:
-- 32
+- 250
 gnn.layers.8.ff_out.linear.weight:
-- 32
-- 32
+- 250
+- 250
 gnn.layers.8.mpnn.model.lin.bias:
-- 32
+- 250
 gnn.layers.8.mpnn.model.lin.weight:
-- 32
+- 250
 - 8
 gnn.layers.8.mpnn.model.nn.fully_connected.0.linear.bias:
-- 32
+- 250
 gnn.layers.8.mpnn.model.nn.fully_connected.0.linear.weight:
-- 32
-- 32
+- 250
+- 250
 gnn.layers.8.mpnn.model.nn.fully_connected.1.linear.bias:
-- 32
+- 250
 gnn.layers.8.mpnn.model.nn.fully_connected.1.linear.weight:
-- 32
-- 32
-gnn.layers.8.norm_layer.bias:
-- 32
-gnn.layers.8.norm_layer.weight:
-- 32
-gnn.layers.8.norm_layer_attn.bias:
-- 32
-gnn.layers.8.norm_layer_attn.weight:
-- 32
-gnn.layers.8.norm_layer_local.bias:
-- 32
-gnn.layers.8.norm_layer_local.weight:
-- 32
+- 250
+- 250
 gnn.layers.9.attn_layer.in_proj_bias:
-- 96
+- 750
 gnn.layers.9.attn_layer.in_proj_weight:
-- 96
-- 32
+- 750
+- 250
 gnn.layers.9.attn_layer.out_proj.bias:
-- 32
+- 250
 gnn.layers.9.attn_layer.out_proj.weight:
-- 32
-- 32
+- 250
+- 250
 gnn.layers.9.ff_linear1.linear.bias:
-- 64
+- 500
 gnn.layers.9.ff_linear1.linear.weight:
-- 64
-- 32
+- 500
+- 250
 gnn.layers.9.ff_linear2.linear.bias:
-- 32
+- 250
 gnn.layers.9.ff_linear2.linear.weight:
-- 32
-- 64
+- 250
+- 500
 gnn.layers.9.ff_out.linear.bias:
-- 32
+- 250
 gnn.layers.9.ff_out.linear.weight:
-- 32
-- 32
+- 250
+- 250
 gnn.layers.9.mpnn.model.lin.bias:
-- 32
+- 250
 gnn.layers.9.mpnn.model.lin.weight:
-- 32
+- 250
 - 8
 gnn.layers.9.mpnn.model.nn.fully_connected.0.linear.bias:
-- 32
+- 250
 gnn.layers.9.mpnn.model.nn.fully_connected.0.linear.weight:
-- 32
-- 32
+- 250
+- 250
 gnn.layers.9.mpnn.model.nn.fully_connected.1.linear.bias:
-- 32
+- 250
 gnn.layers.9.mpnn.model.nn.fully_connected.1.linear.weight:
-- 32
-- 32
-gnn.layers.9.norm_layer.bias:
-- 32
-gnn.layers.9.norm_layer.weight:
-- 32
-gnn.layers.9.norm_layer_attn.bias:
-- 32
-gnn.layers.9.norm_layer_attn.weight:
-- 32
-gnn.layers.9.norm_layer_local.bias:
-- 32
-gnn.layers.9.norm_layer_local.weight:
-- 32
+- 250
+- 250
 gnn.out_linear.linear.bias:
-- 16
+- 250
 gnn.out_linear.linear.weight:
-- 16
-- 48
+- 250
+- 750
 gnn.out_linear.normalization.bias:
-- 16
+- 250
 gnn.out_linear.normalization.weight:
-- 16
+- 250
 pe_encoders.la_pos.linear_A.linear.bias:
 - null
 pe_encoders.la_pos.linear_A.linear.weight:
@@ -861,11 +669,25 @@ pe_encoders.rw_pos.pe_encoder.fully_connected.0.linear.bias:
 pe_encoders.rw_pos.pe_encoder.fully_connected.0.linear.weight:
 - 32
 - null
+post_nn.layers.0.linear.bias:
+- 200
+post_nn.layers.0.linear.weight:
+- 200
+- 250
+post_nn.layers.0.normalization.bias:
+- 200
+post_nn.layers.0.normalization.weight:
+- 200
+post_nn.layers.1.linear.bias:
+- 16
+post_nn.layers.1.linear.weight:
+- 16
+- 200
 pre_nn.layers.0.linear.bias:
 - 16
 pre_nn.layers.0.linear.weight:
 - 16
-- null
+- 87
 pre_nn.layers.0.normalization.bias:
 - 16
 pre_nn.layers.0.normalization.weight:
diff --git a/news/ipu_dispatcher.rst b/news/ipu_dispatcher.rst
new file mode 100644
index 000000000..f72d1171b
--- /dev/null
+++ b/news/ipu_dispatcher.rst
@@ -0,0 +1,11 @@
+**Added:**
+
+* support new compiler based on dispatcher
+* support for latest pytorch-lightning using strategies instead of plugins
+* Simplified IPU installation. Specific IPU requirements file
+* IPU tutorial script along side notebook
+
+**Removed:**
+
+* older IPU compiler based on tracer
+* need to convert custom data structure to tuple of Tensors in order to use the IPU
diff --git a/requirements.txt b/requirements.txt
index cec53248a..fc5fb88e4 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -12,7 +12,7 @@ s3fs >==2021.6
 gcsfs >==2021.6
 appdirs
 tensorboard
-torchmetrics >==0.7.0
+torchmetrics >=0.7.0,<0.11
 hydra-core >==1.0
 ogb
 rdkit >==2020.09
@@ -38,7 +38,7 @@ loguru
 appdirs
 matplotlib
 tk
-pytorch-lightning==1.6
+pytorch-lightning >=1.5.0,<1.9
 datamol
 PyYAML
 wandb
diff --git a/requirements_ipu.txt b/requirements_ipu.txt
new file mode 100644
index 000000000..6c6c109e6
--- /dev/null
+++ b/requirements_ipu.txt
@@ -0,0 +1,11 @@
+--find-links https://data.pyg.org/whl/torch-1.13.0+cpu.html
+--find-links https://download.pytorch.org/whl/cpu/torch_stable.html
+
+pyg-nightly
+torch-scatter==2.1.0
+torch-sparse==0.6.15
+dgl
+torchvision==0.14.0+cpu
+git+https://github.com/joao-alex-cunha/lightning
+
+-r ./requirements.txt
\ No newline at end of file


Test patch:
diff --git a/expts/main_run_test.py b/expts/main_run_test.py
index bcc18a940..acfabc15f 100644
--- a/expts/main_run_test.py
+++ b/expts/main_run_test.py
@@ -23,7 +23,6 @@
 
 
 def main(cfg: DictConfig) -> None:
-
     cfg = deepcopy(cfg)
 
     # Load and initialize the dataset
diff --git a/tests/mtl/test_mtl_architecture.py b/tests/mtl/test_mtl_architecture.py
index f1edae253..29385f4bf 100644
--- a/tests/mtl/test_mtl_architecture.py
+++ b/tests/mtl/test_mtl_architecture.py
@@ -26,17 +26,14 @@
 }
 
 task_1_kwargs = {
-    "task_name": "task_1",
     "out_dim": 5,
     "hidden_dims": [5, 6, 7],
 }
 task_2_kwargs = {
-    "task_name": "task_2",
     "out_dim": 3,
     "hidden_dims": [8, 9, 10],
 }
 task_3_kwargs = {
-    "task_name": "task_3",
     "out_dim": 4,
     "hidden_dims": [2, 2, 2],
 }
@@ -54,11 +51,10 @@
 
 class test_TaskHeads(ut.TestCase):
     def test_task_heads_forward(self):
-
         in_dim = 4  # Dimension of the incoming data
         batch = 2
 
-        task_heads_params = [task_1_params, task_2_params, task_3_params]
+        task_heads_params = {"task_1": task_1_params, "task_2": task_2_params, "task_3": task_3_params}
 
         # Create the "multitask" network. Really it's just an input going to various FFNNs since there's nothing shared.
         multi_head_nn = TaskHeads(in_dim=in_dim, task_heads_kwargs=task_heads_params)
@@ -98,9 +94,9 @@ def test_task_heads_forward(self):
         h_out = multi_head_nn.forward(h)
 
         # Check the output: It's a per-task prediction!
-        self.assertListEqual(list(h_out[task_1_kwargs["task_name"]].shape), [batch, task_1_kwargs["out_dim"]])
-        self.assertListEqual(list(h_out[task_2_kwargs["task_name"]].shape), [batch, task_2_kwargs["out_dim"]])
-        self.assertListEqual(list(h_out[task_3_kwargs["task_name"]].shape), [batch, task_3_kwargs["out_dim"]])
+        self.assertListEqual(list(h_out["task_1"].shape), [batch, task_1_kwargs["out_dim"]])
+        self.assertListEqual(list(h_out["task_2"].shape), [batch, task_2_kwargs["out_dim"]])
+        self.assertListEqual(list(h_out["task_3"].shape), [batch, task_3_kwargs["out_dim"]])
 
 
 class test_Multitask_NN(ut.TestCase):
@@ -143,7 +139,7 @@ class test_Multitask_NN(ut.TestCase):
 
     def test_fulldgl_multitask_forward(self):
         # Task heads
-        task_heads_params = [task_1_params, task_2_params, task_3_params]
+        task_heads_params = {"task_1": task_1_params, "task_2": task_2_params, "task_3": task_3_params}
 
         # Params for the FullDGLNetwork
         temp_dim_1 = 5
@@ -187,24 +183,23 @@ def test_fulldgl_multitask_forward(self):
                             # self.assertListEqual(list(h_out.shape), [dim_1, self.out_dim], msg=err_msg)
 
                             self.assertListEqual(
-                                list(h_out[task_1_kwargs["task_name"]].shape),
+                                list(h_out["task_1"].shape),
                                 [dim_1, task_1_kwargs["out_dim"]],
                                 msg=err_msg,
                             )
                             self.assertListEqual(
-                                list(h_out[task_2_kwargs["task_name"]].shape),
+                                list(h_out["task_2"].shape),
                                 [dim_1, task_2_kwargs["out_dim"]],
                                 msg=err_msg,
                             )
                             self.assertListEqual(
-                                list(h_out[task_3_kwargs["task_name"]].shape),
+                                list(h_out["task_3"].shape),
                                 [dim_1, task_3_kwargs["out_dim"]],
                                 msg=err_msg,
                             )
 
 
 class test_FullGraphMultiTaskNetwork(ut.TestCase):
-
     in_dim_nodes = 7
     in_dim_edges = 13
     in_dims = {"feat": in_dim_nodes, "edge_feat": in_dim_edges}
@@ -233,9 +228,9 @@ def test_FullGraphMultiTaskNetwork_from_config(self):
 
         dim_1 = self.bg.batch_size
 
-        self.assertListEqual(list(h_out[task_1_kwargs["task_name"]].shape), [dim_1, task_1_kwargs["out_dim"]])
-        self.assertListEqual(list(h_out[task_2_kwargs["task_name"]].shape), [dim_1, task_2_kwargs["out_dim"]])
-        self.assertListEqual(list(h_out[task_3_kwargs["task_name"]].shape), [dim_1, task_3_kwargs["out_dim"]])
+        self.assertListEqual(list(h_out["task_1"].shape), [dim_1, task_1_kwargs["out_dim"]])
+        self.assertListEqual(list(h_out["task_2"].shape), [dim_1, task_2_kwargs["out_dim"]])
+        self.assertListEqual(list(h_out["task_3"].shape), [dim_1, task_3_kwargs["out_dim"]])
 
 
 if __name__ == "__main__":
diff --git a/tests/test_architectures.py b/tests/test_architectures.py
index 2bf7ef53b..12d0a044b 100644
--- a/tests/test_architectures.py
+++ b/tests/test_architectures.py
@@ -25,7 +25,6 @@
 
 
 class test_FeedForwardNN(ut.TestCase):
-
     kwargs = {
         "activation": "relu",
         "last_activation": "none",
@@ -321,7 +320,6 @@ def test_forward_weighted_residual_2(self):
 
 
 class test_FeedForwardGraph(ut.TestCase):
-
     kwargs = {
         "activation": "relu",
         "last_activation": "none",
@@ -377,7 +375,6 @@ class test_FeedForwardGraph(ut.TestCase):
     }
 
     def test_forward_no_residual(self):
-
         for pooling in [["none"], ["sum"], ["mean", "logsum", "max"]]:
             for residual_skip_steps in [1, 2]:
                 for virtual_node in self.virtual_nodes:
@@ -431,7 +428,6 @@ def test_forward_no_residual(self):
                             self.assertListEqual(list(h_out.shape), [dim_1, self.out_dim], msg=err_msg)
 
     def test_forward_simple_residual(self):
-
         for pooling in [["none"], ["sum"], ["mean", "logsum", "max"]]:
             for residual_skip_steps in [1, 2]:
                 for virtual_node in self.virtual_nodes:
@@ -484,7 +480,6 @@ def test_forward_simple_residual(self):
                             self.assertListEqual(list(h_out.shape), [dim_1, self.out_dim], msg=err_msg)
 
     def test_forward_weighted_residual(self):
-
         for pooling in [["none"], ["sum"], ["mean", "logsum", "max"]]:
             for residual_skip_steps in [1, 2]:
                 for virtual_node in self.virtual_nodes:
@@ -537,7 +532,6 @@ def test_forward_weighted_residual(self):
                             self.assertListEqual(list(h_out.shape), [dim_1, self.out_dim], msg=err_msg)
 
     def test_forward_concat_residual(self):
-
         for pooling in [["none"], ["sum"], ["mean", "logsum", "max"]]:
             for residual_skip_steps in [1, 2]:
                 for virtual_node in self.virtual_nodes:
@@ -594,7 +588,6 @@ def test_forward_concat_residual(self):
                             self.assertListEqual(list(h_out.shape), [dim_1, self.out_dim], msg=err_msg)
 
     def test_forward_densenet_residual(self):
-
         for pooling in [["none"], ["sum"], ["mean", "logsum", "max"]]:
             for residual_skip_steps in [1, 2]:
                 for virtual_node in self.virtual_nodes:
@@ -654,7 +647,6 @@ def test_forward_densenet_residual(self):
 
 
 class test_FullGraphNetwork(ut.TestCase):
-
     kwargs = {
         "activation": "relu",
         "last_activation": "none",
@@ -711,7 +703,6 @@ class test_FullGraphNetwork(ut.TestCase):
     }
 
     def test_full_network_densenet(self):
-
         temp_dim_1 = 5
         temp_dim_2 = 17
         temp_dim_edges = 21
@@ -733,7 +724,6 @@ def test_full_network_densenet(self):
                             layer_type = layer_name.split("#")[0]
                             for pre_nn_kwargs in pre_nn_kwargs_all:
                                 for pre_nn_edges_kwargs in pre_nn_edges_kwargs_all:
-
                                     # PYG or DGL
                                     if layer_type.startswith("dgl:"):
                                         bg = deepcopy(self.batch_dgl)
diff --git a/tests/test_datamodule.py b/tests/test_datamodule.py
index ba6336bb4..e8408e350 100644
--- a/tests/test_datamodule.py
+++ b/tests/test_datamodule.py
@@ -13,7 +13,6 @@
 class Test_DataModule(ut.TestCase):
     # TODO: Add this test once the OGB Datamodule is fixed
     def test_ogb_datamodule(self):
-
         # other datasets are too large to be tested
         dataset_names = ["ogbg-molhiv", "ogbg-molpcba", "ogbg-moltox21", "ogbg-molfreesolv"]
         dataset_name = dataset_names[3]
@@ -152,7 +151,6 @@ def test_none_filtering(self):
             self.assertIs(series_num_shuffled_2.tolist() == filtered_num, bool_to_check, msg=msg)
 
     def test_caching(self):
-
         # other datasets are too large to be tested
         dataset_name = "ogbg-molfreesolv"
 
diff --git a/tests/test_datasets.py b/tests/test_datasets.py
index d19cfee47..ddb8ad6e8 100644
--- a/tests/test_datasets.py
+++ b/tests/test_datasets.py
@@ -7,7 +7,6 @@
 
 
 class Test_Multitask_Dataset(ut.TestCase):
-
     # Then we can choose different rows and columns for the tests as we see fit.
     # Remember tests are supposed to be FAST, and reading from the file system multiple times slows things down.
 
diff --git a/tests/test_dgl_layers.py b/tests/test_dgl_layers.py
index c4951586f..f3f46414d 100644
--- a/tests/test_dgl_layers.py
+++ b/tests/test_dgl_layers.py
@@ -23,7 +23,6 @@
 
 
 class test_DGL_Layers(ut.TestCase):
-
     in_dim = 21
     out_dim = 11
     in_dim_edges = 13
@@ -46,7 +45,6 @@ class test_DGL_Layers(ut.TestCase):
     }
 
     def test_gcnlayer(self):
-
         bg = deepcopy(self.bg)
         h_in = bg.ndata["h"]
         layer = GCNDgl(in_dim=self.in_dim, out_dim=self.out_dim, **self.kwargs).to(float)
@@ -63,7 +61,6 @@ def test_gcnlayer(self):
         self.assertEqual(h.shape[1], self.out_dim * layer.out_dim_factor)
 
     def test_ginlayer(self):
-
         bg = deepcopy(self.bg)
         h_in = bg.ndata["h"]
         layer = GINDgl(in_dim=self.in_dim, out_dim=self.out_dim, **self.kwargs).to(float)
@@ -80,7 +77,6 @@ def test_ginlayer(self):
         self.assertEqual(h.shape[1], self.out_dim * layer.out_dim_factor)
 
     def test_gatlayer(self):
-
         num_heads = 3
         bg = deepcopy(self.bg)
         h_in = bg.ndata["h"]
@@ -98,7 +94,6 @@ def test_gatlayer(self):
         self.assertEqual(h.shape[1], self.out_dim * layer.out_dim_factor)
 
     def test_gatedgcnlayer(self):
-
         bg = deepcopy(self.bg)
         h_in = bg.ndata["h"]
         e_in = bg.edata["edge_attr"]
@@ -129,7 +124,6 @@ def test_gatedgcnlayer(self):
         self.assertEqual(h.shape[1], self.out_dim * layer.out_dim_factor)
 
     def test_pnaconvolutionallayer(self):
-
         bg = deepcopy(self.bg)
         h_in = bg.ndata["h"]
         e_in = bg.edata["edge_attr"]
@@ -178,7 +172,6 @@ def test_pnaconvolutionallayer(self):
         self.assertEqual(h.shape[1], self.out_dim * layer.out_dim_factor)
 
     def test_pnamessagepassinglayer(self):
-
         bg = deepcopy(self.bg)
         h_in = bg.ndata["h"]
         e_in = bg.edata["edge_attr"]
@@ -228,7 +221,6 @@ def test_pnamessagepassinglayer(self):
         self.assertEqual(h.shape[1], self.out_dim * layer.out_dim_factor)
 
     def test_dgnconvolutionallayer(self):
-
         bg = deepcopy(self.bg)
         h_in = bg.ndata["h"]
         e_in = bg.edata["edge_attr"]
@@ -293,7 +285,6 @@ def test_dgnconvolutionallayer(self):
         self.assertEqual(h.shape[1], self.out_dim * layer.out_dim_factor)
 
     def test_dgnmessagepassinglayer(self):
-
         bg = deepcopy(self.bg)
         h_in = bg.ndata["h"]
         e_in = bg.edata["edge_attr"]
diff --git a/tests/test_featurizer.py b/tests/test_featurizer.py
index 7e30da1d5..aec109790 100644
--- a/tests/test_featurizer.py
+++ b/tests/test_featurizer.py
@@ -20,7 +20,6 @@
 
 
 class test_featurizer(ut.TestCase):
-
     smiles = [
         "C",
         "CC",
@@ -131,7 +130,6 @@ def test_get_mol_atomic_features_float(self):
                 get_mol_atomic_features_float(mol, property_list=bad_props)
 
     def test_get_mol_atomic_features_float_nan_mask(self):
-
         for s in self.smiles_noble:
             mol = dm.to_mol(s)
 
@@ -186,7 +184,6 @@ def test_get_mol_edge_features(self):
                     get_mol_edge_features(mol, property_list=bad_props)
 
     def test_mol_to_adj_and_features(self):
-
         np.random.seed(42)
 
         for s in self.smiles:
@@ -224,7 +221,6 @@ def test_mol_to_adj_and_features(self):
                         self.assertGreaterEqual(ndata.shape[1], num_props, msg=err_msg2)
 
     def test_mol_to_dglgraph(self):
-
         np.random.seed(42)
 
         for s in self.smiles:
@@ -264,7 +260,6 @@ def test_mol_to_dglgraph(self):
                         self.assertGreaterEqual(edata.shape[1], num_props, msg=err_msg2)
 
     def test_mol_to_pyggraph(self):
-
         np.random.seed(42)
 
         for s in self.smiles:
diff --git a/tests/test_ipu_dataloader.py b/tests/test_ipu_dataloader.py
index 1aaa84ad5..280696909 100644
--- a/tests/test_ipu_dataloader.py
+++ b/tests/test_ipu_dataloader.py
@@ -5,7 +5,7 @@
 from copy import deepcopy
 from warnings import warn
 from pytorch_lightning import Trainer, LightningModule
-from pytorch_lightning.plugins import IPUPlugin
+from pytorch_lightning.strategies import IPUStrategy
 from functools import partial
 
 import torch
@@ -13,9 +13,8 @@
 
 # Current library imports
 from goli.ipu.ipu_dataloader import smart_packing, get_pack_sizes, fast_packing, hybrid_packing
-from goli.ipu.ipu_wrapper import PredictorModuleIPU
 from goli.config._loader import load_datamodule, load_metrics, load_architecture
-from goli.ipu.ipu_wrapper import IPUPluginGoli
+from goli.ipu.ipu_utils import import_poptorch
 
 
 def random_packing(num_nodes, batch_size):
@@ -37,7 +36,6 @@ def global_batch_collator(batch_size, batches):
 
 class test_Packing(ut.TestCase):
     def test_smart_packing(self):
-
         np.random.seed(42)
 
         batch_sizes = [2, 4, 8, 16, 32, 64]
@@ -45,7 +43,6 @@ def test_smart_packing(self):
 
         for batch_size in batch_sizes:
             for ipu_batch_size in ipu_batch_sizes:
-
                 err_msg = f"bz={batch_size}, ipu_bz={ipu_batch_size}"
 
                 # Generate random batch size
@@ -78,7 +75,6 @@ def test_smart_packing(self):
                 )
 
     def test_fast_packing(self):
-
         np.random.seed(42)
 
         # Start at 4 for fast_packing for better statistical significance
@@ -87,7 +83,6 @@ def test_fast_packing(self):
 
         for batch_size in batch_sizes:
             for ipu_batch_size in ipu_batch_sizes:
-
                 err_msg = f"bz={batch_size}, ipu_bz={ipu_batch_size}"
 
                 # Generate random batch size
@@ -120,7 +115,6 @@ def test_fast_packing(self):
                 )
 
     def test_hybrid_packing(self):
-
         np.random.seed(42)
 
         batch_sizes = [2, 4, 8, 16, 32, 64]
@@ -128,7 +122,6 @@ def test_hybrid_packing(self):
 
         for batch_size in batch_sizes:
             for ipu_batch_size in ipu_batch_sizes:
-
                 err_msg = f"bz={batch_size}, ipu_bz={ipu_batch_size}"
 
                 # Generate random batch size
@@ -189,107 +182,26 @@ def forward(self, batch):
             return loss
 
         def assert_shapes(self, batch, batch_idx, step):
-            msg = f"\nbatch_idx=`{batch_idx}`, step=`{step}`"
-
             # Test the shape of the labels
             this_shape = list(batch[0].shape)
             true_shape = [1, self.batch_size]
             assert (
                 this_shape == true_shape
-            ), f"Shape of the labels is `{this_shape}` but should be {true_shape}. {msg}"
+            ), f"Shape of the labels is `{this_shape}` but should be {true_shape}"
 
             # Test the shape of the first feature
             this_shape = list(batch[1][0].shape)
             true_shape = [1, self.batch_size, self.node_feat_size]
             assert (
                 this_shape == true_shape
-            ), f"Shape of the feature 0 is `{this_shape}` but should be {true_shape}. {msg}"
+            ), f"Shape of the feature 0 is `{this_shape}` but should be {true_shape}"
 
             # Test the shape of the second feature
             this_shape = list(batch[1][1].shape)
             true_shape = [1, self.batch_size, self.edge_feat_size]
             assert (
                 this_shape == true_shape
-            ), f"Shape of the feature 0 is `{this_shape}` but should be {true_shape}. {msg}"
-
-        def configure_optimizers(self):
-            return torch.optim.Adam(self.parameters(), lr=1e-3)
-
-    class TestPredictor(PredictorModuleIPU):
-        # Create a basic Ligthning for testing the batch sizes
-        def __init__(self, batch_size, node_batch_size, edge_batch_size, in_dims, **kwargs) -> None:
-            super().__init__(**kwargs)
-            self.in_dims = in_dims
-            self.batch_size = batch_size
-            self.node_batch_size = node_batch_size
-            self.edge_batch_size = edge_batch_size
-
-            import poptorch
-
-            self.pop = poptorch
-
-        def validation_step(self, *args):
-            dict_input = self._build_dict_input(*args)
-            self.assert_shapes(args, dict_input, "val")
-            preds = self.forward(dict_input)["preds"]
-            loss = self.compute_loss(
-                preds, dict_input["labels"], weights=None, loss_fun=self.loss_fun, target_nan_mask=None
-            )
-            loss = loss[0]
-            return loss
-
-        def training_step(self, *args):
-            dict_input = self._build_dict_input(*args)
-            self.assert_shapes(args, dict_input, "train")
-            preds = self.forward(dict_input)["preds"]
-            loss = self.compute_loss(
-                preds, dict_input["labels"], weights=None, loss_fun=self.loss_fun, target_nan_mask=None
-            )
-            loss = self.poptorch.identity_loss(loss[0], reduction="mean")
-            return loss
-
-        def assert_shapes(self, args, dict_input, step):
-            msg = f", step=`{step}`"
-
-            # Ensure that the first dimension is 1 for all tensors
-            for arg in args:
-                assert arg.shape[0] == 1
-
-            # Test the shape of the labels
-            this_shape = list(dict_input["labels"]["homo"].shape)
-            true_shape = [self.batch_size, 2]
-            assert (
-                this_shape == true_shape
-            ), f"Shape of the labels is `{this_shape}` but should be {true_shape}. {msg}"
-
-            # Test the shape of the node feature
-            this_shape = list(dict_input["features"]["feat"].shape)
-            true_shape = [self.node_batch_size, self.in_dims["feat"]]
-            assert (
-                this_shape == true_shape
-            ), f"Shape of the feature 0 is `{this_shape}` but should be {true_shape}. {msg}"
-
-            # Test the shape of the node feature
-            this_shape = list(dict_input["features"]["edge_feat"].shape)
-            true_shape = [self.edge_batch_size, self.in_dims["edge_feat"]]
-            assert (
-                this_shape == true_shape
-            ), f"Shape of the feature 0 is `{this_shape}` but should be {true_shape}. {msg}"
-
-        def get_progress_bar_dict(self):
-            return {}
-
-        def on_train_batch_end(self, *args, **kwargs):
-            return
-
-        def on_validation_batch_end(self, *args, **kwargs):
-            return
-
-        def validation_epoch_end(self, *args, **kwargs):
-            return
-
-        def on_train_epoch_end(self) -> None:
-            return
+            ), f"Shape of the feature 0 is `{this_shape}` but should be {true_shape}"
 
         def configure_optimizers(self):
             return torch.optim.Adam(self.parameters(), lr=1e-3)
@@ -332,7 +244,6 @@ def test_poptorch_simple_deviceiterations_gradient_accumulation(self):
         # Initialize the batch info and poptorch options
         opts = poptorch.Options()
         opts.deviceIterations(device_iterations)
-        opts.Jit.traceModel(True)
         training_opts = deepcopy(opts)
         training_opts.Training.gradientAccumulation(gradient_accumulation)
         inference_opts = deepcopy(opts)
@@ -363,9 +274,14 @@ def test_poptorch_simple_deviceiterations_gradient_accumulation(self):
 
         # Build the model, and run it on IPU
         model = self.TestSimpleLightning(batch_size, node_feat_size, edge_feat_size, num_batch)
-        plugins = IPUPlugin(training_opts=training_opts, inference_opts=inference_opts)
+        strategy = IPUStrategy(training_opts=training_opts, inference_opts=inference_opts)
         trainer = Trainer(
-            logger=False, enable_checkpointing=False, max_epochs=2, plugins=plugins, num_sanity_val_steps=0
+            logger=False,
+            enable_checkpointing=False,
+            max_epochs=2,
+            strategy=strategy,
+            num_sanity_val_steps=0,
+            ipus=1,
         )
         trainer.fit(model=model, train_dataloaders=train_dataloader, val_dataloaders=val_dataloader)
 
@@ -382,6 +298,96 @@ def test_poptorch_goli_deviceiterations_gradient_accumulation(self):
             warn(f"Skipping this test because poptorch is not available.\n{e}")
             return
 
+        from goli.ipu.ipu_wrapper import PredictorModuleIPU
+
+        class TestPredictor(PredictorModuleIPU):
+            # Create a basic Ligthning for testing the batch sizes
+            def __init__(self, batch_size, node_batch_size, edge_batch_size, in_dims, **kwargs) -> None:
+                super().__init__(**kwargs)
+                self.in_dims = in_dims
+                self.batch_size = batch_size
+                self.node_batch_size = node_batch_size
+                self.edge_batch_size = edge_batch_size
+
+                import poptorch
+
+                self.pop = poptorch
+
+            def validation_step(self, features, labels):
+                features, labels = self.squeeze_input_dims(features, labels)
+                dict_input = {"features": features, "labels": labels}
+                self.assert_shapes(dict_input, "val")
+                preds = self.forward(dict_input)["preds"]
+                loss = self.compute_loss(
+                    preds, dict_input["labels"], weights=None, loss_fun=self.loss_fun, target_nan_mask=None
+                )
+                loss = loss[0]
+                return loss
+
+            def training_step(self, features, labels):
+                features, labels = self.squeeze_input_dims(features, labels)
+                dict_input = {"features": features, "labels": labels}
+                preds = self.forward(dict_input)["preds"]
+                loss = self.compute_loss(
+                    preds, dict_input["labels"], weights=None, loss_fun=self.loss_fun, target_nan_mask=None
+                )
+                loss = self.pop.identity_loss(loss[0], reduction="mean")
+                return loss
+
+            def assert_shapes(self, dict_input, step):
+                msg = f", step=`{step}`"
+
+                # Test the shape of the labels
+                this_shape = list(dict_input["labels"]["homo"].shape)
+                true_shape = [self.batch_size, 2]
+                assert (
+                    this_shape == true_shape
+                ), f"Shape of the labels is `{this_shape}` but should be {true_shape}. {msg}"
+
+                # Test the shape of the node feature
+                this_shape = list(dict_input["features"]["feat"].shape)
+                true_shape = [self.node_batch_size, self.in_dims["feat"]]
+                assert (
+                    this_shape == true_shape
+                ), f"Shape of the feature 0 is `{this_shape}` but should be {true_shape}. {msg}"
+
+                # Test the shape of the node feature
+                this_shape = list(dict_input["features"]["edge_feat"].shape)
+                true_shape = [self.edge_batch_size, self.in_dims["edge_feat"]]
+                assert (
+                    this_shape == true_shape
+                ), f"Shape of the feature 0 is `{this_shape}` but should be {true_shape}. {msg}"
+
+            def get_progress_bar_dict(self):
+                return {}
+
+            def on_train_batch_end(self, *args, **kwargs):
+                return
+
+            def on_validation_batch_end(self, *args, **kwargs):
+                return
+
+            def validation_epoch_end(self, *args, **kwargs):
+                return
+
+            def on_train_epoch_end(self) -> None:
+                return
+
+            def configure_optimizers(self):
+                return torch.optim.Adam(self.parameters(), lr=1e-3)
+
+            def squeeze_input_dims(self, features, labels):
+                for key, tensor in features:
+                    if isinstance(tensor, torch.Tensor):
+                        features[key] = features[key].squeeze(0)
+
+                for key in labels:
+                    labels[key] = labels[key].squeeze(0)
+
+                return features, labels
+
+        from goli.ipu.ipu_wrapper import DictIPUStrategy
+
         gradient_accumulation = 3
         device_iterations = 5
         batch_size = 7
@@ -389,7 +395,6 @@ def test_poptorch_goli_deviceiterations_gradient_accumulation(self):
         # Initialize the batch info and poptorch options
         opts = poptorch.Options()
         opts.deviceIterations(device_iterations)
-        opts.Jit.traceModel(True)
         training_opts = deepcopy(opts)
         training_opts.Training.gradientAccumulation(gradient_accumulation)
         inference_opts = deepcopy(opts)
@@ -412,7 +417,7 @@ def test_poptorch_goli_deviceiterations_gradient_accumulation(self):
 
         model_class, model_kwargs = load_architecture(cfg, in_dims=datamodule.in_dims)
         metrics = load_metrics(cfg)
-        predictor = self.TestPredictor(
+        predictor = TestPredictor(
             batch_size=batch_size,
             node_batch_size=node_factor * batch_size,
             edge_batch_size=edge_factor * batch_size,
@@ -422,9 +427,9 @@ def test_poptorch_goli_deviceiterations_gradient_accumulation(self):
             metrics=metrics,
             **cfg["predictor"],
         )
-        plugins = IPUPluginGoli(training_opts=training_opts, inference_opts=inference_opts)
+        strategy = DictIPUStrategy(training_opts=training_opts, inference_opts=inference_opts)
         trainer = Trainer(
-            logger=False, enable_checkpointing=False, max_epochs=2, plugins=plugins, num_sanity_val_steps=0
+            logger=False, enable_checkpointing=False, max_epochs=2, strategy=strategy, num_sanity_val_steps=0
         )
         trainer.fit(model=predictor, datamodule=datamodule)
 
diff --git a/tests/test_ipu_losses.py b/tests/test_ipu_losses.py
index db5e0dec4..05bcd3e8e 100644
--- a/tests/test_ipu_losses.py
+++ b/tests/test_ipu_losses.py
@@ -7,7 +7,6 @@
 
 
 class test_Losses(ut.TestCase):
-
     torch.manual_seed(42)
     preds = torch.rand((100, 10), dtype=torch.float32)
     target = torch.rand((100, 10), dtype=torch.float32)
diff --git a/tests/test_ipu_metrics.py b/tests/test_ipu_metrics.py
index 660df8127..dbaaab8ef 100644
--- a/tests/test_ipu_metrics.py
+++ b/tests/test_ipu_metrics.py
@@ -33,7 +33,6 @@
 
 
 class test_Metrics(ut.TestCase):
-
     torch.manual_seed(42)
     preds = torch.rand((100, 10), dtype=torch.float32)
     target = torch.rand((100, 10), dtype=torch.float32)
diff --git a/tests/test_metrics.py b/tests/test_metrics.py
index c99050167..e4a5e5fdf 100644
--- a/tests/test_metrics.py
+++ b/tests/test_metrics.py
@@ -18,7 +18,6 @@
 
 class test_Metrics(ut.TestCase):
     def test_thresholder(self):
-
         torch.manual_seed(42)
         preds = torch.rand(100, dtype=torch.float32)
         target = torch.rand(100, dtype=torch.float32)
@@ -62,11 +61,9 @@ def test_thresholder(self):
 
 class test_MetricWrapper(ut.TestCase):
     def test_target_nan_mask(self):
-
         torch.random.manual_seed(42)
 
         for sz in [(100,), (100, 1), (100, 10)]:
-
             err_msg = f"Error for `sz = {sz}`"
 
             # Generate prediction and target matrices
diff --git a/tests/test_multitask_datamodule.py b/tests/test_multitask_datamodule.py
index 7c2fa3133..4ce50a57e 100644
--- a/tests/test_multitask_datamodule.py
+++ b/tests/test_multitask_datamodule.py
@@ -117,7 +117,6 @@ def test_multitask_fromsmiles_dm(self):
             assert batch["labels"]["score"].shape == (16, 1)
 
     def test_multitask_fromsmiles_from_config(self):
-
         config = goli.load_config(name="zinc_default_multitask_fulldgl")
 
         df = goli.data.load_tiny_zinc()  # 100 molecules
diff --git a/tests/test_mup.py b/tests/test_mup.py
index 7b156cdf3..24c29fa0a 100644
--- a/tests/test_mup.py
+++ b/tests/test_mup.py
@@ -136,7 +136,6 @@ def test_feedforwardgraph_mup(self):
         model_2_lastreadout.forward(deepcopy(in_features))
 
     def test_fullgraphnetwork(self):
-
         # Load the configuration file for the model
         CONFIG_FILE = "tests/config_test_ipu_dataloader.yaml"
         with open(CONFIG_FILE, "r") as f:
@@ -223,7 +222,6 @@ def test_fullgraphnetwork(self):
         model_2.forward(deepcopy(in_features))
 
     def test_fullgraphmultitasknetwork(self):
-
         # Load the configuration file for the model
         CONFIG_FILE = "tests/config_test_ipu_dataloader.yaml"
         with open(CONFIG_FILE, "r") as f:
diff --git a/tests/test_positional_encoding.py b/tests/test_positional_encoding.py
index 97b77bdf4..667c8f3ec 100644
--- a/tests/test_positional_encoding.py
+++ b/tests/test_positional_encoding.py
@@ -16,7 +16,6 @@
 
 
 class test_positional_encoder(ut.TestCase):
-
     smiles = [
         "C",
         "CC",
@@ -29,7 +28,6 @@ class test_positional_encoder(ut.TestCase):
     adjs = [Chem.rdmolops.GetAdjacencyMatrix(mol) for mol in mols]
 
     def test_laplacian_eigvec(self):
-
         for ii, adj in enumerate(deepcopy(self.adjs)):
             for num_pos in [1, 2, 4]:  # Can't test too much eigs because of multiplicities
                 for disconnected_comp in [True, False]:
@@ -70,7 +68,6 @@ def test_laplacian_eigvec(self):
                         )
 
     def test_laplacian_eigvec_eigval(self):
-
         for ii, adj in enumerate(deepcopy(self.adjs)):
             for num_pos in [1, 2, 4]:  # Can't test too much eigs because of multiplicities
                 for disconnected_comp in [True, False]:
@@ -130,7 +127,6 @@ def test_rwse(self):
     """
 
     def test_laplacian_eigvec_with_encoder(self):
-
         for ii, adj in enumerate(deepcopy(self.adjs)):
             for num_pos in [2, 4, 8]:  # Can't test too much eigs because of multiplicities
                 for disconnected_comp in [True, False]:
diff --git a/tests/test_pyg_layers.py b/tests/test_pyg_layers.py
index 41c706313..84a0e99a2 100644
--- a/tests/test_pyg_layers.py
+++ b/tests/test_pyg_layers.py
@@ -22,7 +22,6 @@
 
 
 class test_Pyg_Layers(ut.TestCase):
-
     in_dim = 21
     out_dim = 11
     in_dim_edges = 13
@@ -68,7 +67,6 @@ def test_gpslayer(self):
         self.assertEqual(bg.h.shape[1], self.out_dim * layer.out_dim_factor)
 
     def test_ginlayer(self):
-
         bg = deepcopy(self.bg)
         h_in = bg.h
         layer = GINConvPyg(in_dim=self.in_dim, out_dim=self.out_dim, **self.kwargs)
@@ -85,7 +83,6 @@ def test_ginlayer(self):
         self.assertEqual(bg.h.shape[1], self.out_dim * layer.out_dim_factor)
 
     def test_ginelayer(self):
-
         bg = deepcopy(self.bg)
         h_in = bg.h
         layer = GINEConvPyg(in_dim=self.in_dim, out_dim=self.out_dim, **self.kwargs)
@@ -107,7 +104,6 @@ def test_ginelayer(self):
         self.assertEqual(bg.h.shape[1], self.out_dim * layer.out_dim_factor)
 
     def test_gatedgcnlayer(self):
-
         bg = deepcopy(self.bg)
         h_in = bg.h
         e_in = bg.edge_attr
@@ -134,7 +130,6 @@ def test_gatedgcnlayer(self):
         self.assertEqual(bg2.h.shape[1], self.out_dim * layer.out_dim_factor)
 
     def test_pnamessagepassinglayer(self):
-
         bg = deepcopy(self.bg)
         h_in = bg.h
         aggregators = ["mean", "max", "min", "std", "sum"]
diff --git a/tests/test_residual_connections.py b/tests/test_residual_connections.py
index 068041ec9..f25c18347 100644
--- a/tests/test_residual_connections.py
+++ b/tests/test_residual_connections.py
@@ -50,7 +50,6 @@ def test_get_true_out_dims_simple(self):
         self.assertListEqual(expected_out_dims, true_out_dims)
 
     def test_forward_simple(self):
-
         for skip_steps in [1, 2, 3]:
             rc = ResidualConnectionSimple(skip_steps=skip_steps)
             num_loops = 10
@@ -99,7 +98,6 @@ def test_get_true_out_dims_random(self):
             rc = ResidualConnectionRandom(skip_steps=1, out_dims=None, num_layers=None)
 
     def test_forward_random(self):
-
         for skip_steps in [1, 2, 3]:
             num_loops = 10
 
@@ -130,7 +128,6 @@ def test_get_true_out_dims_weighted(self):
         self.assertListEqual(expected_out_dims, true_out_dims)
 
     def test_forward_weighted(self):
-
         for skip_steps in [1, 2, 3]:
             num_loops = 10
             shape = (3, 11)
@@ -149,7 +146,6 @@ def test_forward_weighted(self):
             h_prev = None
             step_counter = 0
             for ii in range(num_loops):
-
                 h_prev_backup = h_prev
                 h, h_prev = rc.forward(h_original[ii], h_prev, step_idx=ii)
 
@@ -197,7 +193,6 @@ def test_get_true_out_dims_concat(self):
         self.assertListEqual(expected_out_dims, true_out_dims)
 
     def test_forward_concat(self):
-
         for skip_steps in [1, 2, 3]:
             rc = ResidualConnectionConcat(skip_steps=skip_steps)
             num_loops = 10
@@ -250,7 +245,6 @@ def test_get_true_out_dims_densenet(self):
         self.assertListEqual(expected_out_dims, true_out_dims)
 
     def test_forward_densenet(self):
-
         for skip_steps in [1, 2, 3]:
             rc = ResidualConnectionDenseNet(skip_steps=skip_steps)
             num_loops = 10
diff --git a/tests/test_utils.py b/tests/test_utils.py
index 10e538c28..c5ba7f555 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -11,7 +11,6 @@
 
 
 class test_nan_statistics(ut.TestCase):
-
     torch.manual_seed(42)
 
     dims = [
@@ -32,7 +31,6 @@ class test_nan_statistics(ut.TestCase):
     tensor[is_nan] = float("nan")
 
     def test_nan_mean(self):
-
         for keepdim in [False, True]:
             for dim in self.dims:
                 err_msg = f"Error for :\n dim = {dim}\n keepdim = {keepdim}"
@@ -53,7 +51,6 @@ def test_nan_mean(self):
                 np.testing.assert_almost_equal(torch_mean.numpy(), numpy_mean, decimal=6, err_msg=err_msg)
 
     def test_nan_std_var(self):
-
         for unbiased in [True, False]:
             for keepdim in [False, True]:
                 for dim in self.dims:
@@ -80,7 +77,6 @@ def test_nan_std_var(self):
                     np.testing.assert_almost_equal(torch_var.numpy(), numpy_var, decimal=6, err_msg=err_msg)
 
     def test_nan_median(self):
-
         for keepdim in [False, True]:
             # Cannot test
             for dim in self.dims:  # in [d for d in self.dims if not isinstance(d, Tuple)]:
@@ -106,7 +102,6 @@ def test_nan_median(self):
                 self.assertListEqual(list(torch_med.shape), list(torch_sum.shape))
 
     def test_nan_mad(self):
-
         for normal in [False, True]:
             # Cannot test
             for dim in self.dims:  # in [d for d in self.dims if not isinstance(d, Tuple)]:
@@ -132,7 +127,6 @@ def test_nan_mad(self):
 
 class test_SafeRun(ut.TestCase):
     def test_safe_run(self):
-
         # Error is caught
         with SafeRun(name="bob", raise_error=False, verbose=0):
             raise ValueError("This is an error")