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

Issue:
Remove ipu
## Changelogs

- Removing support for IPU
- Removing custom `to_dense_batch` and `to_sparse_batch` specific to IPUs, and reverted to the pyg version
- Removing all documentation regarding IPUs

---

_discussion related to that PR_



Discussion / hints:

Fix patch:
diff --git a/.gitignore b/.gitignore
index 2afd2115e..751e6aeb9 100644
--- a/.gitignore
+++ b/.gitignore
@@ -55,15 +55,6 @@ debug/
 change_commits.sh
 graphium/features/test_new_pes.ipynb
 
-# IPU related ignores and profiler outputs
-*.a
-*.cbor
-*.capnp
-*.pop
-*.popart
-*.pop_cache
-*.popef
-*.pvti*
 
 ############ END graphium Custom GitIgnore ##############
 
diff --git a/README.md b/README.md
index 53a7172bb..7556b0a65 100644
--- a/README.md
+++ b/README.md
@@ -13,7 +13,6 @@
 [![GitHub Repo stars](https://img.shields.io/github/stars/datamol-io/graphium)](https://github.com/datamol-io/graphium/stargazers)
 [![GitHub Repo stars](https://img.shields.io/github/forks/datamol-io/graphium)](https://github.com/datamol-io/graphium/network/members)
 [![test](https://github.com/datamol-io/graphium/actions/workflows/test.yml/badge.svg)](https://github.com/datamol-io/graphium/actions/workflows/test.yml)
-[![test-ipu](https://github.com/datamol-io/graphium/actions/workflows/test_ipu.yml/badge.svg)](https://github.com/datamol-io/graphium/actions/workflows/test_ipu.yml)
 [![release](https://github.com/datamol-io/graphium/actions/workflows/release.yml/badge.svg)](https://github.com/datamol-io/graphium/actions/workflows/release.yml)
 [![code-check](https://github.com/datamol-io/graphium/actions/workflows/code-check.yml/badge.svg)](https://github.com/datamol-io/graphium/actions/workflows/code-check.yml)
 [![doc](https://github.com/datamol-io/graphium/actions/workflows/doc.yml/badge.svg)](https://github.com/datamol-io/graphium/actions/workflows/doc.yml)
@@ -53,18 +52,6 @@ mamba activate graphium
 pip install --no-deps -e .
 ```
 
-### For IPU developers
-```bash
-# Install Graphcore's SDK and Graphium dependencies in a new environment called `.graphium_ipu`
-./install_ipu.sh .graphium_ipu
-```
-
-The above step needs to be done once. After that, enable the SDK and the environment as follows:
-
-```bash
-source enable_ipu.sh .graphium_ipu
-```
-
 ## Training a model
 
 To learn how to train a model, we invite you to look at the documentation, or the jupyter notebooks available [here](https://github.com/datamol-io/graphium/tree/master/docs/tutorials/model_training).
diff --git a/codecov.yml b/codecov.yml
index 60a2c37d6..94e8bd149 100644
--- a/codecov.yml
+++ b/codecov.yml
@@ -16,8 +16,3 @@ component_management:
         target: auto
         branches:
           - "!main"
-  individual_components:
-    - component_id: ipu # this is an identifier that should not be changed
-      name: ipu # this is a display name, and can be changed freely
-      paths:
-        - graphium/ipu/**
diff --git a/docs/api/graphium.ipu.md b/docs/api/graphium.ipu.md
deleted file mode 100644
index 3943e78a4..000000000
--- a/docs/api/graphium.ipu.md
+++ /dev/null
@@ -1,43 +0,0 @@
-graphium.ipu
-====================
-Code for adapting to run on IPU
-
-=== "Contents"
-
-    * [IPU Dataloader](#ipu-dataloader)
-    * [IPU Losses](#ipu-losses)
-    * [IPU Metrics](#ipu-metrics)
-    * [IPU Simple Lightning](#ipu-simple-lightning)
-    * [IPU Utils](#ipu-utils)
-    * [IPU Wrapper](#ipu-wrapper)
-    * [To Dense Batch](#to-dense-batch)
-
-## IPU Dataloader
-------------
-::: graphium.ipu.ipu_dataloader
-
-
-## IPU Losses
-------------
-::: graphium.ipu.ipu_losses
-
-
-## IPU Simple Lightning
-------------
-::: graphium.ipu.ipu_simple_lightning
-
-
-## IPU Utils
-------------
-::: graphium.ipu.ipu_utils
-
-
-## IPU Wrapper
-------------
-::: graphium.ipu.ipu_wrapper
-
-
-## To Dense Batch
-------------
-::: graphium.ipu.to_dense_batch
-
diff --git a/docs/contribute.md b/docs/contribute.md
index b4fef7ce0..4f9f71fae 100644
--- a/docs/contribute.md
+++ b/docs/contribute.md
@@ -18,21 +18,6 @@ mamba activate graphium
 pip install --no-deps -e .
 ```
 
-### For IPU developers
-
-Download the SDK and use pypi to create your environment:
-
-```bash
-# Install Graphcore's SDK and Graphium dependencies in a new environment called `.graphium_ipu`
-./install_ipu.sh .graphium_ipu
-```
-
-The above step needs to be done once. After that, enable the SDK and the environment as follows:
-
-```bash
-source enable_ipu.sh .graphium_ipu
-```
-
 ## Build the documentation
 
 You can build and serve the documentation locally with:
diff --git a/docs/design.md b/docs/design.md
index 380ac28e4..43594fc98 100644
--- a/docs/design.md
+++ b/docs/design.md
@@ -42,7 +42,6 @@ Below are a list of directory and their respective documentations:
 - [data](https://github.com/datamol-io/graphium/blob/main/graphium/data/README.md)
 - [features](https://github.com/datamol-io/graphium/tree/main/graphium/features/README.md)
 - finetuning
-- [ipu](https://github.com/datamol-io/graphium/tree/main/graphium/ipu/README.md)
 - [nn](https://github.com/datamol-io/graphium/tree/main/graphium/nn/README.md)
 - [trainer](https://github.com/datamol-io/graphium/tree/main/graphium/trainer/README.md)
 - [utils](https://github.com/datamol-io/graphium/tree/main/graphium/features/README.md)
@@ -56,7 +55,7 @@ Hence, we use [hydra](https://hydra.cc/docs/intro/) to enable splitting the conf
 
 Examples of possibilities include:
 
-- Switching between accelerators (CPU, GPU and IPU)
+- Switching between accelerators (CPU, GPU)
 - Benchmarking different models on the same dataset
 - Fine-tuning a pre-trained model on a new dataset
 
diff --git a/docs/index.md b/docs/index.md
index ef01e54be..32170ee29 100644
--- a/docs/index.md
+++ b/docs/index.md
@@ -25,17 +25,6 @@ or pip:
 pip install graphium
 ```
 
-### For IPU
-```bash
-# Install Graphcore's SDK and Graphium dependencies in a new environment called `.graphium_ipu`
-./install_ipu.sh .graphium_ipu
-```
-
-The above step needs to be done once. After that, enable the SDK and the environment as follows:
-
-```bash
-source enable_ipu.sh .graphium_ipu
-```
 
 Finally, you will need to install graphium with pip
 ```bash
diff --git a/docs/tutorials/feature_processing/timing_parallel.ipynb b/docs/tutorials/feature_processing/timing_parallel.ipynb
index 477251e71..22e3552e5 100644
--- a/docs/tutorials/feature_processing/timing_parallel.ipynb
+++ b/docs/tutorials/feature_processing/timing_parallel.ipynb
@@ -476,7 +476,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "graphium_ipu",
+   "display_name": "graphium",
    "language": "python",
    "name": "python3"
   },
diff --git a/docs/tutorials/model_training/simple-molecular-model.ipynb b/docs/tutorials/model_training/simple-molecular-model.ipynb
index 26a45cfa0..545109ebc 100644
--- a/docs/tutorials/model_training/simple-molecular-model.ipynb
+++ b/docs/tutorials/model_training/simple-molecular-model.ipynb
@@ -405,9 +405,9 @@
       "    zinc:\n",
       "    - mae\n",
       "  loss_fun:\n",
-      "    qm9: mae_ipu\n",
-      "    tox21: bce_logits_ipu\n",
-      "    zinc: mae_ipu\n",
+      "    qm9: mae\n",
+      "    tox21: bce_logits\n",
+      "    zinc: mae\n",
       "  random_seed: ${constants.seed}\n",
       "  optim_kwargs:\n",
       "    lr: 4.0e-05\n",
@@ -451,28 +451,28 @@
       "metrics:\n",
       "  qm9:\n",
       "  - name: mae\n",
-      "    metric: mae_ipu\n",
+      "    metric: mae\n",
       "    target_nan_mask: null\n",
       "    multitask_handling: flatten\n",
       "    threshold_kwargs: null\n",
       "  - name: pearsonr\n",
-      "    metric: pearsonr_ipu\n",
+      "    metric: pearsonr\n",
       "    threshold_kwargs: null\n",
       "    target_nan_mask: null\n",
       "    multitask_handling: mean-per-label\n",
       "  - name: r2_score\n",
-      "    metric: r2_score_ipu\n",
+      "    metric: r2_score\n",
       "    target_nan_mask: null\n",
       "    multitask_handling: mean-per-label\n",
       "    threshold_kwargs: null\n",
       "  tox21:\n",
       "  - name: auroc\n",
-      "    metric: auroc_ipu\n",
+      "    metric: auroc\n",
       "    task: binary\n",
       "    multitask_handling: mean-per-label\n",
       "    threshold_kwargs: null\n",
       "  - name: avpr\n",
-      "    metric: average_precision_ipu\n",
+      "    metric: average_precision\n",
       "    task: binary\n",
       "    multitask_handling: mean-per-label\n",
       "    threshold_kwargs: null\n",
@@ -498,17 +498,17 @@
       "      th_on_target: true\n",
       "  zinc:\n",
       "  - name: mae\n",
-      "    metric: mae_ipu\n",
+      "    metric: mae\n",
       "    target_nan_mask: null\n",
       "    multitask_handling: flatten\n",
       "    threshold_kwargs: null\n",
       "  - name: pearsonr\n",
-      "    metric: pearsonr_ipu\n",
+      "    metric: pearsonr\n",
       "    threshold_kwargs: null\n",
       "    target_nan_mask: null\n",
       "    multitask_handling: mean-per-label\n",
       "  - name: r2_score\n",
-      "    metric: r2_score_ipu\n",
+      "    metric: r2_score\n",
       "    target_nan_mask: null\n",
       "    multitask_handling: mean-per-label\n",
       "    threshold_kwargs: null\n",
diff --git a/enable_ipu.sh b/enable_ipu.sh
deleted file mode 100755
index 63dd34987..000000000
--- a/enable_ipu.sh
+++ /dev/null
@@ -1,29 +0,0 @@
-"""
---------------------------------------------------------------------------------
-Copyright (c) 2023 Graphcore Limited.
-Use of this software is subject to the terms and conditions outlined in the LICENSE file.
-Unauthorized modification, distribution, or use is prohibited. Provided 'as is' without
-warranties of any kind.
-
-Graphcore Limited is not liable
-for any damages arising from its use.
-Refer to the LICENSE file for the full terms and conditions.
---------------------------------------------------------------------------------
-"""
-
-
-#!/bin/bash
-
-# Default location for the virtual environment
-default_venv_name=".graphium_ipu"
-
-# Allow the user to specify the location of their virtual environment
-# If not specified, use the default location
-venv_name=${1:-$default_venv_name}
-
-# Constants
-sdk_path="${venv_name}/poplar_sdk-ubuntu_20_04-3.3.0+1403-208993bbb7"
-
-# Source the virtual environment
-source ${venv_name}/bin/activate
-source ${sdk_path}/enable
\ No newline at end of file
diff --git a/env.yml b/env.yml
index d3a3cbc23..f73a5d219 100644
--- a/env.yml
+++ b/env.yml
@@ -70,7 +70,16 @@ dependencies:
   - markdown-include
   - mike >=1.0.0
 
+  # C++ dependencies
+  - gcc_linux-64 # Sometimes I find that I need to enforce `gcc_linux-64`, but that won't work with Mac, Windows, or Arm-Linux
+  - gxx_linux-64 # Sometimes I find that I need to enforce `gxx_linux-64`, but that won't work with Mac, Windows, or Arm-Linux
+  - libgcc
+  - pybind11
+  - boost
+
+  # Optional
+  - pytdc
+
   - pip:
-      - lightning-graphcore # optional, for using IPUs only
       - hydra-core>=1.3.2
       - hydra-optuna-sweeper
diff --git a/expts/configs/config_gps_10M_pcqm4m.yaml b/expts/configs/config_gps_10M_pcqm4m.yaml
index 0487a8d04..8fc38812e 100644
--- a/expts/configs/config_gps_10M_pcqm4m.yaml
+++ b/expts/configs/config_gps_10M_pcqm4m.yaml
@@ -1,22 +1,15 @@
-# Testing the mpnn only model with the PCQMv2 dataset on IPU.
+# Testing the mpnn only model with the PCQMv2 dataset.
 constants:
   name: &name pcqm4mv2_mpnn_4layer
   seed: &seed 42
   raise_train_error: true   # Whether the code should raise an error if it crashes during training
 
 accelerator:
-  type: ipu  # cpu or ipu or gpu
+  type: gpu  # cpu or gpu
   config_override:
     datamodule:
       args:
-        ipu_dataloader_training_opts:
-          mode: async
-          max_num_nodes_per_graph: 20 # train max nodes: 20, max_edges: 54
-          max_num_edges_per_graph: 60
-        ipu_dataloader_inference_opts:
-          mode: async
-          max_num_nodes_per_graph: 16 # valid max nodes: 51, max_edges: 118
-          max_num_edges_per_graph: 120
+
         # Data handling-related
         batch_size_training: 64
         batch_size_inference: 16
@@ -28,35 +21,6 @@ accelerator:
         precision: 16
         accumulate_grad_batches: 4
 
-  ipu_config:
-    - deviceIterations(20) # IPU would require large batches to be ready for the model.
-    - replicationFactor(16)
-    # - enableProfiling("graph_analyser")       # The folder where the profile will be stored
-    # - enableExecutableCaching("pop_compiler_cache")
-    - TensorLocations.numIOTiles(128)
-    - _Popart.set("defaultBufferingDepth", 128)
-    - Precision.enableStochasticRounding(True)
-
-  ipu_inference_config: # Optional. If not provided, same as `ipu_config`
-    - deviceIterations(80) # IPU would require large batches to be ready for the model.
-    - replicationFactor(16)
-    # - enableProfiling("graph_analyser")       # The folder where the profile will be stored
-    # - enableExecutableCaching("pop_compiler_cache")
-    - TensorLocations.numIOTiles(128)
-    - _Popart.set("defaultBufferingDepth", 128)
-    - Precision.enableStochasticRounding(True)
-
-# accelerator:
-#   type: cpu  # cpu or ipu or gpu
-#   config_override:
-#     datamodule:
-#       batch_size_training: 256
-#       batch_size_inference: 64
-#     trainer:
-#       trainer:
-#         precision: 32
-#         accumulate_grad_batches: 1
-
 datamodule:
   module_type: "MultitaskFromSmilesDataModule"
   args: # Matches that in the test_multitask_datamodule.py case.
@@ -212,7 +176,7 @@ predictor:
   metrics_on_progress_bar:
     homolumo: ["mae", "pearsonr"]
   loss_fun:
-    homolumo: mse_ipu
+    homolumo: mse
   random_seed: *seed
   optim_kwargs:
     lr: 4.e-4 # warmup can be scheduled using torch_scheduler_kwargs
@@ -235,12 +199,12 @@ predictor:
 metrics:
   homolumo:
     - name: mae
-      metric: mae_ipu
+      metric: mae
       target_nan_mask: null
       multitask_handling: flatten
       threshold_kwargs: null
     - name: pearsonr
-      metric: pearsonr_ipu
+      metric: pearsonr
       threshold_kwargs: null
       target_nan_mask: null
       multitask_handling: mean-per-label
diff --git a/expts/configs/config_gps_10M_pcqm4m_mod.yaml b/expts/configs/config_gps_10M_pcqm4m_mod.yaml
index 19543302b..1587d2d06 100644
--- a/expts/configs/config_gps_10M_pcqm4m_mod.yaml
+++ b/expts/configs/config_gps_10M_pcqm4m_mod.yaml
@@ -1,10 +1,10 @@
-# Testing the mpnn only model with the PCQMv2 dataset on IPU.
+# Testing the mpnn only model with the PCQMv2 dataset.
 constants:
   name: &name pcqm4mv2_mpnn_4layer
   seed: &seed 42
   raise_train_error: true   # Whether the code should raise an error if it crashes during training
   accelerator:
-    type: gpu  # cpu or ipu or gpu
+    type: gpu  # cpu or gpu
 
 datamodule:
   module_type: "MultitaskFromSmilesDataModule"
@@ -80,18 +80,6 @@ datamodule:
     persistent_workers: False # if use persistent worker at the start of each epoch.
     # Using persistent_workers false might make the start of each epoch very long.
 
-    # ipu_dataloader_training_opts:
-    #   mode: async
-    #   max_num_nodes_per_graph: 20 # train max nodes: 20, max_edges: 54
-    #   max_num_edges_per_graph: 60
-
-    # ipu_dataloader_inference_opts:
-    #   mode: async
-    #   max_num_nodes_per_graph: 20 # valid max nodes: 51, max_edges: 118
-    #   max_num_edges_per_graph: 120
-    #   # test-dev max nodes: 50, max_edges: 116
-    #   # test-challenge max nodes: 51, max_edges: 106
-
 architecture:
   model_type: FullGraphMultiTaskNetwork
   mup_base_path: null
@@ -223,7 +211,7 @@ predictor:
   metrics_on_progress_bar:
     homolumo: ["mae", "pearsonr"]
   loss_fun:
-    homolumo: mse_ipu
+    homolumo: mse
   random_seed: *seed
   optim_kwargs:
     lr: 4.e-4 # warmup can be scheduled using torch_scheduler_kwargs
@@ -246,12 +234,12 @@ predictor:
 metrics:
   homolumo:
     - name: mae
-      metric: mae_ipu
+      metric: mae
       target_nan_mask: null
       multitask_handling: flatten
       threshold_kwargs: null
     - name: pearsonr
-      metric: pearsonr_ipu
+      metric: pearsonr
       threshold_kwargs: null
       target_nan_mask: null
       multitask_handling: mean-per-label
diff --git a/expts/configs/config_mpnn_10M_b3lyp.yaml b/expts/configs/config_mpnn_10M_b3lyp.yaml
index 424dbcd71..8403a1ba1 100644
--- a/expts/configs/config_mpnn_10M_b3lyp.yaml
+++ b/expts/configs/config_mpnn_10M_b3lyp.yaml
@@ -1,22 +1,15 @@
-# Testing the mpnn only model with the b3lyp dataset on IPU.
+# Testing the mpnn only model with the b3lyp dataset.
 constants:
   name: &name b3lyp_mpnn_4layer
   seed: &seed 42
   raise_train_error: true   # Whether the code should raise an error if it crashes during training
 
 accelerator:
-  type: ipu  # cpu or ipu or gpu
+  type: gpu  # cpu or gpu
   config_override:
     datamodule:
       args:
-        ipu_dataloader_training_opts:
-          mode: async
-          max_num_nodes_per_graph: 20 # train max nodes: 20, max_edges: 54
-          max_num_edges_per_graph: 60
-        ipu_dataloader_inference_opts:
-          mode: async
-          max_num_nodes_per_graph: 16 # valid max nodes: 51, max_edges: 118
-          max_num_edges_per_graph: 120
+
         # Data handling-related
         batch_size_training: 64
         batch_size_inference: 16
@@ -28,36 +21,6 @@ accelerator:
         precision: 16
         accumulate_grad_batches: 4
 
-  ipu_config:
-    - deviceIterations(20) # IPU would require large batches to be ready for the model.
-    - replicationFactor(16)
-    # - enableProfiling("graph_analyser")       # The folder where the profile will be stored
-    # - enableExecutableCaching("pop_compiler_cache")
-    - TensorLocations.numIOTiles(128)
-    - _Popart.set("defaultBufferingDepth", 128)
-    - Precision.enableStochasticRounding(True)
-
-  ipu_inference_config: # Optional. If not provided, same as `ipu_config`
-    - deviceIterations(80) # IPU would require large batches to be ready for the model.
-    - replicationFactor(16)
-    # - enableProfiling("graph_analyser")       # The folder where the profile will be stored
-    # - enableExecutableCaching("pop_compiler_cache")
-    - TensorLocations.numIOTiles(128)
-    - _Popart.set("defaultBufferingDepth", 128)
-    - Precision.enableStochasticRounding(True)
-
-
-# accelerator:
-#   type: cpu  # cpu or ipu or gpu
-#   config_override:
-#     datamodule:
-#       batch_size_training: 256
-#       batch_size_inference: 64
-#     trainer:
-#       trainer:
-#         precision: 32
-#         accumulate_grad_batches: 1
-
 datamodule:
   module_type: "MultitaskFromSmilesDataModule"
   args: # Matches that in the test_multitask_datamodule.py case.
@@ -244,8 +207,8 @@ predictor:
     alphagap: ["mae", "pearsonr"]
     betagap: ["mae", "pearsonr"]
   loss_fun:
-    alphagap: mse_ipu
-    betagap: mse_ipu
+    alphagap: mse
+    betagap: mse
   random_seed: *seed
   optim_kwargs:
     lr: 4.e-4 # warmup can be scheduled using torch_scheduler_kwargs
@@ -268,12 +231,12 @@ predictor:
 metrics:
   alphagap: &alpha_metrics
     - name: mae
-      metric: mae_ipu
+      metric: mae
       target_nan_mask: null
       multitask_handling: flatten
       threshold_kwargs: null
     - name: pearsonr
-      metric: pearsonr_ipu
+      metric: pearsonr
       threshold_kwargs: null
       target_nan_mask: null
       multitask_handling: mean-per-label
diff --git a/expts/configs/config_mpnn_pcqm4m.yaml b/expts/configs/config_mpnn_pcqm4m.yaml
index 70972d370..0ba8d3bd5 100644
--- a/expts/configs/config_mpnn_pcqm4m.yaml
+++ b/expts/configs/config_mpnn_pcqm4m.yaml
@@ -1,10 +1,10 @@
-# Testing the mpnn only model with the PCQMv2 dataset on IPU.
+# Testing the mpnn only model with the PCQMv2 dataset.
 constants:
   name: &name pcqm4mv2_mpnn_4layer
   seed: &seed 42
   raise_train_error: true   # Whether the code should raise an error if it crashes during training
   accelerator:
-    type: cpu  # cpu or ipu or gpu
+    type: cpu  # cpu or gpu
 
 datamodule:
   module_type: "MultitaskFromSmilesDataModule"
@@ -56,18 +56,6 @@ datamodule:
     persistent_workers: False # if use persistent worker at the start of each epoch.
     # Using persistent_workers false might make the start of each epoch very long.
 
-    # ipu_dataloader_training_opts:
-    #   mode: async
-    #   max_num_nodes_per_graph: 20 # train max nodes: 20, max_edges: 54
-    #   max_num_edges_per_graph: 60
-
-    # ipu_dataloader_inference_opts:
-    #   mode: async
-    #   max_num_nodes_per_graph: 20 # valid max nodes: 51, max_edges: 118
-    #   max_num_edges_per_graph: 120
-    #   # test-dev max nodes: 50, max_edges: 116
-    #   # test-challenge max nodes: 51, max_edges: 106
-
 architecture:
   model_type: FullGraphMultiTaskNetwork
   mup_base_path: null
@@ -166,7 +154,7 @@ predictor:
   metrics_on_progress_bar:
     homolumo: ["mae", "pearsonr"]
   loss_fun:
-    homolumo: mse_ipu
+    homolumo: mse
   random_seed: *seed
   optim_kwargs:
     lr: 4.e-4 # warmup can be scheduled using torch_scheduler_kwargs
@@ -189,12 +177,12 @@ predictor:
 metrics:
   homolumo:
     - name: mae
-      metric: mae_ipu
+      metric: mae
       target_nan_mask: null
       multitask_handling: flatten
       threshold_kwargs: null
     - name: pearsonr
-      metric: pearsonr_ipu
+      metric: pearsonr
       threshold_kwargs: null
       target_nan_mask: null
       multitask_handling: mean-per-label
diff --git a/expts/dataset_benchmark.py b/expts/dataset_benchmark.py
index e8bf2e24d..948a03688 100644
--- a/expts/dataset_benchmark.py
+++ b/expts/dataset_benchmark.py
@@ -21,7 +21,6 @@
 # CONFIG_FILE = "expts/neurips2023_configs/debug/config_large_gcn_debug.yaml"
 CONFIG_FILE = "expts/neurips2023_configs/config_large_gcn.yaml"
 # CONFIG_FILE = "expts/configs/config_pcqmv2_mpnn.yaml"
-# CONFIG_FILE = "expts/configs/config_ipu_qm9.yaml"
 
 
 def benchmark(fn, *args, message="", log2wandb=False, **kwargs):
diff --git a/expts/hydra-configs/README.md b/expts/hydra-configs/README.md
index 40625917d..77d2569ba 100644
--- a/expts/hydra-configs/README.md
+++ b/expts/hydra-configs/README.md
@@ -1,14 +1,14 @@
 # Configuring Graphium with Hydra
 This document provides users with a point of entry to composing configs in Graphium. As a flexible library with many features, configuration is an important part of Graphium. To make configurations as reusable as possible while providing maximum flexibility, we integrated Graphium with `hydra`. Our config structure is designed to make the following functionality as accessible as possible:
 
-- Switching between **accelerators** (CPU, GPU and IPU)
+- Switching between **accelerators** (CPU, GPU)
 - **Benchmarking** different models on the same dataset
 - **Fine-tuning** a pre-trained model on a new dataset
 
 In what follows, we describe how each of the above functionality is achieved and how users can benefit from this design to achieve the most with Graphium with as little configuration as possible.
 
 ## Accelerators
-With Graphium supporting CPU, GPU and IPU hardware, easily switching between these accelerators is pre-configured. General, accelerator-specific configs are specified under `accelerator/`, whereas experiment-specific differences between the accelerators are specialized under `training/accelerator`.
+With Graphium supporting CPU, GPU hardware, easily switching between these accelerators is pre-configured. General, accelerator-specific configs are specified under `accelerator/`, whereas experiment-specific differences between the accelerators are specialized under `training/accelerator`.
 
 ## Benchmarking
 Benchmarking multiple models on the same datasets and tasks requires us to easily switch between model configurations without redefining major parts of the architecture, task heads, featurization, metrics, predictor, etc. For example, when changing from a GCN to a GIN model, a simple switch of `architecture.gnn.layer_type: 'pyg:gin'` might suffice. Hence, we abstract the `model` configs under `model/` where such model configurations can be specified.
diff --git a/expts/hydra-configs/accelerator/ipu.yaml b/expts/hydra-configs/accelerator/ipu.yaml
deleted file mode 100644
index 3e6fb4429..000000000
--- a/expts/hydra-configs/accelerator/ipu.yaml
+++ /dev/null
@@ -1,18 +0,0 @@
-type: ipu
-ipu_config:
-    - deviceIterations(60) # IPU would require large batches to be ready for the model.
-    # 60 for PCQM4mv2
-    # 30 for largemix
-    - replicationFactor(16)
-    # - enableProfiling("graph_analyser")       # The folder where the profile will be stored
-    # - enableExecutableCaching("pop_compiler_cache")
-    - TensorLocations.numIOTiles(128)
-    - _Popart.set("defaultBufferingDepth", 96)
-    - Precision.enableStochasticRounding(True)
-
-ipu_inference_config:
-    # set device iteration and replication factor to 1 during inference
-    # gradient accumulation was set to 1 in the code
-    - deviceIterations(1)
-    - replicationFactor(1)
-    - Precision.enableStochasticRounding(False)
diff --git a/expts/hydra-configs/accelerator/ipu_pipeline.yaml b/expts/hydra-configs/accelerator/ipu_pipeline.yaml
deleted file mode 100644
index 996218646..000000000
--- a/expts/hydra-configs/accelerator/ipu_pipeline.yaml
+++ /dev/null
@@ -1,22 +0,0 @@
-type: ipu
-ipu_config:
-    - deviceIterations(60) # IPU would require large batches to be ready for the model.
-    # 60 for PCQM4mv2
-    # 30 for largemix
-    - replicationFactor(4)
-    # - enableProfiling("graph_analyser")       # The folder where the profile will be stored
-    # - enableExecutableCaching("pop_compiler_cache")
-    - TensorLocations.numIOTiles(128)
-    - _Popart.set("defaultBufferingDepth", 96)
-    - Precision.enableStochasticRounding(True)
-
-ipu_inference_config:
-    # set device iteration and replication factor to 1 during inference
-    # gradient accumulation was set to 1 in the code
-    - deviceIterations(60)
-    - replicationFactor(1)
-    - Precision.enableStochasticRounding(False)
-
-accelerator_kwargs:
-    _accelerator: "ipu"
-    gnn_layers_per_ipu: [4, 4, 4, 4]
\ No newline at end of file
diff --git a/expts/hydra-configs/tasks/loss_metrics_datamodule/l1000_mcf7.yaml b/expts/hydra-configs/tasks/loss_metrics_datamodule/l1000_mcf7.yaml
index 43933a7fa..53c753402 100644
--- a/expts/hydra-configs/tasks/loss_metrics_datamodule/l1000_mcf7.yaml
+++ b/expts/hydra-configs/tasks/loss_metrics_datamodule/l1000_mcf7.yaml
@@ -7,7 +7,7 @@ predictor:
     l1000_mcf7: []
   loss_fun:
     l1000_mcf7:
-      name: hybrid_ce_ipu
+      name: hybrid_ce
       n_brackets: 3
       alpha: 0.5
 
diff --git a/expts/hydra-configs/tasks/loss_metrics_datamodule/l1000_vcap.yaml b/expts/hydra-configs/tasks/loss_metrics_datamodule/l1000_vcap.yaml
index 27b89d862..e385bf23e 100644
--- a/expts/hydra-configs/tasks/loss_metrics_datamodule/l1000_vcap.yaml
+++ b/expts/hydra-configs/tasks/loss_metrics_datamodule/l1000_vcap.yaml
@@ -7,7 +7,7 @@ predictor:
     l1000_vcap: []
   loss_fun:
     l1000_vcap:
-      name: hybrid_ce_ipu
+      name: hybrid_ce
       n_brackets: 3
       alpha: 0.5
 
diff --git a/expts/hydra-configs/tasks/loss_metrics_datamodule/largemix.yaml b/expts/hydra-configs/tasks/loss_metrics_datamodule/largemix.yaml
index 921960cd1..f307f0441 100644
--- a/expts/hydra-configs/tasks/loss_metrics_datamodule/largemix.yaml
+++ b/expts/hydra-configs/tasks/loss_metrics_datamodule/largemix.yaml
@@ -15,16 +15,16 @@ predictor:
     pcqm4m_n4: []
   loss_fun:
     l1000_vcap:
-      name: hybrid_ce_ipu
+      name: hybrid_ce
       n_brackets: 3
       alpha: 0.5
     l1000_mcf7:
-      name: hybrid_ce_ipu
+      name: hybrid_ce
       n_brackets: 3
       alpha: ${predictor.loss_fun.l1000_vcap.alpha}
-    pcba_1328: bce_logits_ipu
-    pcqm4m_g25: mae_ipu
-    pcqm4m_n4: mae_ipu
+    pcba_1328: bce_logits
+    pcqm4m_g25: mae
+    pcqm4m_n4: mae
 
 metrics:
   l1000_vcap: &classif_metrics
@@ -48,7 +48,7 @@ metrics:
       threshold_kwargs: null
   l1000_mcf7: *classif_metrics
   pcba_1328:
-  # use auroc and averageprecision (non_ipu version) so tha nans are handled correctly
+  
     - name: auroc
       metric: auroc
       task: binary
@@ -63,17 +63,17 @@ metrics:
       threshold_kwargs: null
   pcqm4m_g25: &pcqm_metrics
     - name: mae
-      metric: mae_ipu
+      metric: mae
       target_nan_mask: null
       multitask_handling: mean-per-label
       threshold_kwargs: null
     - name: pearsonr
-      metric: pearsonr_ipu
+      metric: pearsonr
       threshold_kwargs: null
       target_nan_mask: null
       multitask_handling: mean-per-label
     - name: r2
-      metric: r2_score_ipu
+      metric: r2_score
       threshold_kwargs: null
       target_nan_mask: null
       multitask_handling: mean-per-label
diff --git a/expts/hydra-configs/tasks/loss_metrics_datamodule/pcba_1328.yaml b/expts/hydra-configs/tasks/loss_metrics_datamodule/pcba_1328.yaml
index adc3321a0..72f9fba35 100644
--- a/expts/hydra-configs/tasks/loss_metrics_datamodule/pcba_1328.yaml
+++ b/expts/hydra-configs/tasks/loss_metrics_datamodule/pcba_1328.yaml
@@ -6,11 +6,11 @@ predictor:
   metrics_on_training_set:
     pcba_1328: []
   loss_fun:
-    pcba_1328: bce_logits_ipu
+    pcba_1328: bce_logits
 
 metrics:
   pcba_1328:
-  # use auroc and averageprecision (non_ipu version) so tha nans are handled correctly
+  
     - name: auroc
       metric: auroc
       task: binary
diff --git a/expts/hydra-configs/tasks/loss_metrics_datamodule/pcqm4m.yaml b/expts/hydra-configs/tasks/loss_metrics_datamodule/pcqm4m.yaml
index 8dcf2c0c4..8eb878b62 100644
--- a/expts/hydra-configs/tasks/loss_metrics_datamodule/pcqm4m.yaml
+++ b/expts/hydra-configs/tasks/loss_metrics_datamodule/pcqm4m.yaml
@@ -7,18 +7,18 @@ predictor:
   metrics_on_training_set:
     homolumo: ["pearsonr"]  
   loss_fun:
-    homolumo: mae_ipu
+    homolumo: mae
 
 # Task-specific
 metrics:
   homolumo:
     - name: mae
-      metric: mae_ipu
+      metric: mae
       target_nan_mask: null
       multitask_handling: mean-per-label
       threshold_kwargs: null
     - name: pearsonr
-      metric: pearsonr_ipu
+      metric: pearsonr
       threshold_kwargs: null
       target_nan_mask: null
       multitask_handling: mean-per-label
diff --git a/expts/hydra-configs/tasks/loss_metrics_datamodule/pcqm4m_g25.yaml b/expts/hydra-configs/tasks/loss_metrics_datamodule/pcqm4m_g25.yaml
index 047701f6e..8247b4c47 100644
--- a/expts/hydra-configs/tasks/loss_metrics_datamodule/pcqm4m_g25.yaml
+++ b/expts/hydra-configs/tasks/loss_metrics_datamodule/pcqm4m_g25.yaml
@@ -6,22 +6,22 @@ predictor:
   metrics_on_training_set:
     pcqm4m_g25: []
   loss_fun:
-    pcqm4m_g25: mae_ipu
+    pcqm4m_g25: mae
 
 metrics:
   pcqm4m_g25:
     - name: mae
-      metric: mae_ipu
+      metric: mae
       target_nan_mask: null
       multitask_handling: mean-per-label
       threshold_kwargs: null
     - name: pearsonr
-      metric: pearsonr_ipu
+      metric: pearsonr
       threshold_kwargs: null
       target_nan_mask: null
       multitask_handling: mean-per-label
     - name: r2
-      metric: r2_score_ipu
+      metric: r2_score
       threshold_kwargs: null
       target_nan_mask: null
       multitask_handling: mean-per-label
diff --git a/expts/hydra-configs/tasks/loss_metrics_datamodule/pcqm4m_n4.yaml b/expts/hydra-configs/tasks/loss_metrics_datamodule/pcqm4m_n4.yaml
index 494843464..2ef471be4 100644
--- a/expts/hydra-configs/tasks/loss_metrics_datamodule/pcqm4m_n4.yaml
+++ b/expts/hydra-configs/tasks/loss_metrics_datamodule/pcqm4m_n4.yaml
@@ -6,22 +6,22 @@ predictor:
   metrics_on_training_set:
     pcqm4m_n4: []
   loss_fun:
-    pcqm4m_n4: mae_ipu
+    pcqm4m_n4: mae
 
 metrics:
   pcqm4m_n4:
     - name: mae
-      metric: mae_ipu
+      metric: mae
       target_nan_mask: null
       multitask_handling: mean-per-label
       threshold_kwargs: null
     - name: pearsonr
-      metric: pearsonr_ipu
+      metric: pearsonr
       threshold_kwargs: null
       target_nan_mask: null
       multitask_handling: mean-per-label
     - name: r2
-      metric: r2_score_ipu
+      metric: r2_score
       threshold_kwargs: null
       target_nan_mask: null
       multitask_handling: mean-per-label
diff --git a/expts/hydra-configs/tasks/loss_metrics_datamodule/toymix.yaml b/expts/hydra-configs/tasks/loss_metrics_datamodule/toymix.yaml
index ceb4e6d69..bf2e044b4 100644
--- a/expts/hydra-configs/tasks/loss_metrics_datamodule/toymix.yaml
+++ b/expts/hydra-configs/tasks/loss_metrics_datamodule/toymix.yaml
@@ -8,9 +8,9 @@ predictor:
     tox21: ["auroc"]
     zinc: ["mae"]
   loss_fun:
-    qm9: mae_ipu
-    tox21: bce_logits_ipu
-    zinc: mae_ipu
+    qm9: mae
+    tox21: bce_logits
+    zinc: mae
 
 metrics:
   qm9: &qm9_metrics
diff --git a/expts/hydra-configs/training/accelerator/largemix_ipu.yaml b/expts/hydra-configs/training/accelerator/largemix_ipu.yaml
index 090600e98..115cd9e53 100644
--- a/expts/hydra-configs/training/accelerator/largemix_ipu.yaml
+++ b/expts/hydra-configs/training/accelerator/largemix_ipu.yaml
@@ -2,14 +2,6 @@
 
 datamodule:
     args:
-      ipu_dataloader_training_opts:
-        mode: async
-        max_num_nodes_per_graph: 30 # train max nodes: 20, max_edges: 54
-        max_num_edges_per_graph: 100
-      ipu_dataloader_inference_opts:
-        mode: async
-        max_num_nodes_per_graph: 35 # valid max nodes: 51, max_edges: 118
-        max_num_edges_per_graph: 100
       # Data handling-related
       batch_size_training: 30
       batch_size_inference: 30
diff --git a/expts/hydra-configs/training/accelerator/pcqm4m_ipu.yaml b/expts/hydra-configs/training/accelerator/pcqm4m_ipu.yaml
index a7e23f383..c49d10405 100644
--- a/expts/hydra-configs/training/accelerator/pcqm4m_ipu.yaml
+++ b/expts/hydra-configs/training/accelerator/pcqm4m_ipu.yaml
@@ -2,14 +2,6 @@
 
 datamodule:
   args:
-    ipu_dataloader_training_opts:
-      mode: async
-      max_num_nodes_per_graph: 16 # train max nodes: 20, max_edges: 54
-      max_num_edges_per_graph: 60
-    ipu_dataloader_inference_opts:
-      mode: async
-      max_num_nodes_per_graph: 30 # valid max nodes: 51, max_edges: 118
-      max_num_edges_per_graph: 120
     # Data handling-related
     batch_size_inference: 16
 
diff --git a/expts/hydra-configs/training/accelerator/toymix_ipu.yaml b/expts/hydra-configs/training/accelerator/toymix_ipu.yaml
index 1bf28ce0b..8f5fe4941 100644
--- a/expts/hydra-configs/training/accelerator/toymix_ipu.yaml
+++ b/expts/hydra-configs/training/accelerator/toymix_ipu.yaml
@@ -2,14 +2,6 @@
 
 datamodule:
   args:
-    ipu_dataloader_training_opts:
-      mode: async
-      max_num_nodes_per_graph: 44 # train max nodes: 20, max_edges: 54
-      max_num_edges_per_graph: 80
-    ipu_dataloader_inference_opts:
-      mode: async
-      max_num_nodes_per_graph: 44 # valid max nodes: 51, max_edges: 118
-      max_num_edges_per_graph: 80
     # Data handling-related
     batch_size_training: 50
     batch_size_inference: 50
diff --git a/expts/hydra-configs/training/toymix.yaml b/expts/hydra-configs/training/toymix.yaml
index dc6c174b9..bfffdebe2 100644
--- a/expts/hydra-configs/training/toymix.yaml
+++ b/expts/hydra-configs/training/toymix.yaml
@@ -11,7 +11,7 @@ predictor:
     warmup_epochs: 10
     verbose: False
   scheduler_kwargs: null
-  target_nan_mask: null
+  target_nan_mask: ignore
   multitask_handling: flatten # flatten, mean-per-label
 
 trainer:
diff --git a/expts/neurips2023_configs/base_config/large.yaml b/expts/neurips2023_configs/base_config/large.yaml
index 18503527c..748680247 100644
--- a/expts/neurips2023_configs/base_config/large.yaml
+++ b/expts/neurips2023_configs/base_config/large.yaml
@@ -7,18 +7,11 @@ constants:
   datacache_path: "/localdata/neurips2023-large/"
 
 accelerator:
-  type: ipu  # cpu or ipu or gpu
+  type: gpu  # cpu or gpu
   config_override:
     datamodule:
       args:
-        ipu_dataloader_training_opts:
-          mode: async
-          max_num_nodes_per_graph: 30 # train max nodes: 20, max_edges: 54
-          max_num_edges_per_graph: 100
-        ipu_dataloader_inference_opts:
-          mode: async
-          max_num_nodes_per_graph: 35 # valid max nodes: 51, max_edges: 118
-          max_num_edges_per_graph: 100
+
         # Data handling-related
         batch_size_training: 30
         batch_size_inference: 30
@@ -31,35 +24,6 @@ accelerator:
         precision: 16-true
         accumulate_grad_batches: 2
 
-  ipu_config:
-    - deviceIterations(30) # IPU would require large batches to be ready for the model.
-    - replicationFactor(16)
-    # - enableProfiling("graph_analyser")       # The folder where the profile will be stored
-    # - enableExecutableCaching("pop_compiler_cache")
-    - TensorLocations.numIOTiles(128)
-    - _Popart.set("defaultBufferingDepth", 96)
-    - Precision.enableStochasticRounding(True)
-    # - Precision.enableFloatingPointExceptions(True)
-
-  ipu_inference_config:
-  # set device iteration and replication factor to 1 during inference
-  # gradient accumulation was set to 1 in the code
-    - deviceIterations(1)
-    - replicationFactor(1)
-    - Precision.enableStochasticRounding(False)
-
-# accelerator:
-#   type: cpu  # cpu or ipu or gpu
-#   config_override:
-#     datamodule:
-#       args:
-#         batch_size_training: 64
-#         batch_size_inference: 256
-#     trainer:
-#       trainer:
-#         precision: 32
-#         accumulate_grad_batches: 1
-
 datamodule:
   module_type: "MultitaskFromSmilesDataModule"
   args: # Matches that in the test_multitask_datamodule.py case.
@@ -325,16 +289,16 @@ predictor:
     pcqm4m_n4: []
   loss_fun:
     l1000_vcap:
-      name: hybrid_ce_ipu
+      name: hybrid_ce
       n_brackets: 3
       alpha: 0.5
     l1000_mcf7:
-      name: hybrid_ce_ipu
+      name: hybrid_ce
       n_brackets: 3
       alpha: ${predictor.loss_fun.l1000_vcap.alpha}
-    pcba_1328: bce_logits_ipu
-    pcqm4m_g25: mae_ipu
-    pcqm4m_n4: mae_ipu
+    pcba_1328: bce_logits
+    pcqm4m_g25: mae
+    pcqm4m_n4: mae
   random_seed: ${constants.seed}
   optim_kwargs:
     lr: 1.e-4 # warmup can be scheduled using torch_scheduler_kwargs
@@ -374,7 +338,7 @@ metrics:
       threshold_kwargs: null
   l1000_mcf7: *classif_metrics
   pcba_1328:
-  # use auroc and averageprecision (non_ipu version) so tha nans are handled correctly
+  
     - name: auroc
       metric: auroc
       task: binary
@@ -389,17 +353,17 @@ metrics:
       threshold_kwargs: null
   pcqm4m_g25: &pcqm_metrics
     - name: mae
-      metric: mae_ipu
+      metric: mae
       target_nan_mask: null
       multitask_handling: mean-per-label
       threshold_kwargs: null
     - name: pearsonr
-      metric: pearsonr_ipu
+      metric: pearsonr
       threshold_kwargs: null
       target_nan_mask: null
       multitask_handling: mean-per-label
     - name: r2
-      metric: r2_score_ipu
+      metric: r2_score
       threshold_kwargs: null
       target_nan_mask: null
       multitask_handling: mean-per-label
diff --git a/expts/neurips2023_configs/base_config/large_pcba.yaml b/expts/neurips2023_configs/base_config/large_pcba.yaml
index a1e3d108f..ca016e27d 100644
--- a/expts/neurips2023_configs/base_config/large_pcba.yaml
+++ b/expts/neurips2023_configs/base_config/large_pcba.yaml
@@ -7,18 +7,11 @@ constants:
   datacache_path: "/localdata/neurips2023-large/"
 
 accelerator:
-  type: ipu  # cpu or ipu or gpu
+  type: gpu  # cpu or gpu
   config_override:
     datamodule:
       args:
-        ipu_dataloader_training_opts:
-          mode: async
-          max_num_nodes_per_graph: 30 # train max nodes: 20, max_edges: 54
-          max_num_edges_per_graph: 100
-        ipu_dataloader_inference_opts:
-          mode: async
-          max_num_nodes_per_graph: 35 # valid max nodes: 51, max_edges: 118
-          max_num_edges_per_graph: 100
+
         # Data handling-related
         batch_size_training: 30
         batch_size_inference: 30
@@ -31,35 +24,6 @@ accelerator:
         precision: 16-true
         accumulate_grad_batches: 2
 
-  ipu_config:
-    - deviceIterations(30) # IPU would require large batches to be ready for the model.
-    - replicationFactor(16)
-    # - enableProfiling("graph_analyser")       # The folder where the profile will be stored
-    # - enableExecutableCaching("pop_compiler_cache")
-    - TensorLocations.numIOTiles(128)
-    - _Popart.set("defaultBufferingDepth", 96)
-    - Precision.enableStochasticRounding(True)
-    # - Precision.enableFloatingPointExceptions(True)
-
-  ipu_inference_config:
-  # set device iteration and replication factor to 1 during inference
-  # gradient accumulation was set to 1 in the code
-    - deviceIterations(1)
-    - replicationFactor(1)
-    - Precision.enableStochasticRounding(False)
-
-# accelerator:
-#   type: cpu  # cpu or ipu or gpu
-#   config_override:
-#     datamodule:
-#       args:
-#         batch_size_training: 64
-#         batch_size_inference: 256
-#     trainer:
-#       trainer:
-#         precision: 32
-#         accumulate_grad_batches: 1
-
 datamodule:
   module_type: "MultitaskFromSmilesDataModule"
   args: # Matches that in the test_multitask_datamodule.py case.
@@ -324,16 +288,16 @@ predictor:
     #pcqm4m_n4: []
   loss_fun:
     # l1000_vcap:
-    #   name: hybrid_ce_ipu
+    #   name: hybrid_ce
     #   n_brackets: 3
     #   alpha: 0.5
     # l1000_mcf7:
-    #   name: hybrid_ce_ipu
+    #   name: hybrid_ce
     #   n_brackets: 3
     #   alpha: ${predictor.loss_fun.l1000_vcap.alpha}
-    pcba_1328: bce_logits_ipu
-    # pcqm4m_g25: mae_ipu
-    #pcqm4m_n4: mae_ipu
+    pcba_1328: bce_logits
+    # pcqm4m_g25: mae
+    #pcqm4m_n4: mae
   random_seed: ${constants.seed}
   optim_kwargs:
     lr: 1.e-4 # warmup can be scheduled using torch_scheduler_kwargs
@@ -373,7 +337,7 @@ metrics:
   #     threshold_kwargs: null
   # l1000_mcf7: *classif_metrics
   pcba_1328:
-  # use auroc and averageprecision (non_ipu version) so tha nans are handled correctly
+  
     - name: auroc
       metric: auroc
       task: binary
@@ -388,17 +352,17 @@ metrics:
       threshold_kwargs: null
       # pcqm4m_n4: &pcqm_metrics
       #- name: mae
-      #metric: mae_ipu
+      #metric: mae
       #target_nan_mask: null
       #multitask_handling: mean-per-label
       #threshold_kwargs: null
       #- name: pearsonr
-      #metric: pearsonr_ipu
+      #metric: pearsonr
       #threshold_kwargs: null
       #target_nan_mask: null
       #multitask_handling: mean-per-label
       #- name: r2
-      #metric: r2_score_ipu
+      #metric: r2_score
       #threshold_kwargs: null
       #target_nan_mask: null
       #multitask_handling: mean-per-label
diff --git a/expts/neurips2023_configs/base_config/large_pcqm_g25.yaml b/expts/neurips2023_configs/base_config/large_pcqm_g25.yaml
index b71c43cf2..12b6afb2d 100644
--- a/expts/neurips2023_configs/base_config/large_pcqm_g25.yaml
+++ b/expts/neurips2023_configs/base_config/large_pcqm_g25.yaml
@@ -7,18 +7,11 @@ constants:
   datacache_path: "/localdata/neurips2023-large/"
 
 accelerator:
-  type: ipu  # cpu or ipu or gpu
+  type: gpu  # cpu or gpu
   config_override:
     datamodule:
       args:
-        ipu_dataloader_training_opts:
-          mode: async
-          max_num_nodes_per_graph: 30 # train max nodes: 20, max_edges: 54
-          max_num_edges_per_graph: 100
-        ipu_dataloader_inference_opts:
-          mode: async
-          max_num_nodes_per_graph: 35 # valid max nodes: 51, max_edges: 118
-          max_num_edges_per_graph: 100
+
         # Data handling-related
         batch_size_training: 30
         batch_size_inference: 30
@@ -31,35 +24,6 @@ accelerator:
         precision: 16-true
         accumulate_grad_batches: 2
 
-  ipu_config:
-    - deviceIterations(30) # IPU would require large batches to be ready for the model.
-    - replicationFactor(16)
-    # - enableProfiling("graph_analyser")       # The folder where the profile will be stored
-    # - enableExecutableCaching("pop_compiler_cache")
-    - TensorLocations.numIOTiles(128)
-    - _Popart.set("defaultBufferingDepth", 96)
-    - Precision.enableStochasticRounding(True)
-    # - Precision.enableFloatingPointExceptions(True)
-
-  ipu_inference_config:
-  # set device iteration and replication factor to 1 during inference
-  # gradient accumulation was set to 1 in the code
-    - deviceIterations(1)
-    - replicationFactor(1)
-    - Precision.enableStochasticRounding(False)
-
-# accelerator:
-#   type: cpu  # cpu or ipu or gpu
-#   config_override:
-#     datamodule:
-#       args:
-#         batch_size_training: 64
-#         batch_size_inference: 256
-#     trainer:
-#       trainer:
-#         precision: 32
-#         accumulate_grad_batches: 1
-
 datamodule:
   module_type: "MultitaskFromSmilesDataModule"
   args: # Matches that in the test_multitask_datamodule.py case.
@@ -324,16 +288,16 @@ predictor:
     # pcqm4m_n4: []
   loss_fun:
     # l1000_vcap:
-    #   name: hybrid_ce_ipu
+    #   name: hybrid_ce
     #   n_brackets: 3
     #   alpha: 0.5
     # l1000_mcf7:
-    #   name: hybrid_ce_ipu
+    #   name: hybrid_ce
     #   n_brackets: 3
     #   alpha: ${predictor.loss_fun.l1000_vcap.alpha}
-    # pcba_1328: bce_logits_ipu
-    pcqm4m_g25: mae_ipu
-    # pcqm4m_n4: mae_ipu
+    # pcba_1328: bce_logits
+    pcqm4m_g25: mae
+    # pcqm4m_n4: mae
   random_seed: ${constants.seed}
   optim_kwargs:
     lr: 1.e-4 # warmup can be scheduled using torch_scheduler_kwargs
@@ -373,7 +337,7 @@ metrics:
   #     threshold_kwargs: null
   # l1000_mcf7: *classif_metrics
   # pcba_1328:
-  # # use auroc and averageprecision (non_ipu version) so tha nans are handled correctly
+  # 
   #   - name: auroc
   #     metric: auroc
   #     task: binary
@@ -388,17 +352,17 @@ metrics:
   #     threshold_kwargs: null
   pcqm4m_g25: &pcqm_metrics
     - name: mae
-      metric: mae_ipu
+      metric: mae
       target_nan_mask: null
       multitask_handling: mean-per-label
       threshold_kwargs: null
     - name: pearsonr
-      metric: pearsonr_ipu
+      metric: pearsonr
       threshold_kwargs: null
       target_nan_mask: null
       multitask_handling: mean-per-label
     - name: r2
-      metric: r2_score_ipu
+      metric: r2_score
       threshold_kwargs: null
       target_nan_mask: null
       multitask_handling: mean-per-label
diff --git a/expts/neurips2023_configs/base_config/large_pcqm_n4.yaml b/expts/neurips2023_configs/base_config/large_pcqm_n4.yaml
index 464e49581..5741854dc 100644
--- a/expts/neurips2023_configs/base_config/large_pcqm_n4.yaml
+++ b/expts/neurips2023_configs/base_config/large_pcqm_n4.yaml
@@ -7,18 +7,11 @@ constants:
   datacache_path: "/localdata/neurips2023-large/"
 
 accelerator:
-  type: ipu  # cpu or ipu or gpu
+  type: gpu  # cpu or gpu
   config_override:
     datamodule:
       args:
-        ipu_dataloader_training_opts:
-          mode: async
-          max_num_nodes_per_graph: 30 # train max nodes: 20, max_edges: 54
-          max_num_edges_per_graph: 100
-        ipu_dataloader_inference_opts:
-          mode: async
-          max_num_nodes_per_graph: 35 # valid max nodes: 51, max_edges: 118
-          max_num_edges_per_graph: 100
+
         # Data handling-related
         batch_size_training: 30
         batch_size_inference: 30
@@ -31,35 +24,6 @@ accelerator:
         precision: 16-true
         accumulate_grad_batches: 2
 
-  ipu_config:
-    - deviceIterations(30) # IPU would require large batches to be ready for the model.
-    - replicationFactor(16)
-    # - enableProfiling("graph_analyser")       # The folder where the profile will be stored
-    # - enableExecutableCaching("pop_compiler_cache")
-    - TensorLocations.numIOTiles(128)
-    - _Popart.set("defaultBufferingDepth", 96)
-    - Precision.enableStochasticRounding(True)
-    # - Precision.enableFloatingPointExceptions(True)
-
-  ipu_inference_config:
-  # set device iteration and replication factor to 1 during inference
-  # gradient accumulation was set to 1 in the code
-    - deviceIterations(1)
-    - replicationFactor(1)
-    - Precision.enableStochasticRounding(False)
-
-# accelerator:
-#   type: cpu  # cpu or ipu or gpu
-#   config_override:
-#     datamodule:
-#       args:
-#         batch_size_training: 64
-#         batch_size_inference: 256
-#     trainer:
-#       trainer:
-#         precision: 32
-#         accumulate_grad_batches: 1
-
 datamodule:
   module_type: "MultitaskFromSmilesDataModule"
   args: # Matches that in the test_multitask_datamodule.py case.
@@ -324,16 +288,16 @@ predictor:
     pcqm4m_n4: []
   loss_fun:
     # l1000_vcap:
-    #   name: hybrid_ce_ipu
+    #   name: hybrid_ce
     #   n_brackets: 3
     #   alpha: 0.5
     # l1000_mcf7:
-    #   name: hybrid_ce_ipu
+    #   name: hybrid_ce
     #   n_brackets: 3
     #   alpha: ${predictor.loss_fun.l1000_vcap.alpha}
-    # pcba_1328: bce_logits_ipu
-    # pcqm4m_g25: mae_ipu
-    pcqm4m_n4: mae_ipu
+    # pcba_1328: bce_logits
+    # pcqm4m_g25: mae
+    pcqm4m_n4: mae
   random_seed: ${constants.seed}
   optim_kwargs:
     lr: 1.e-4 # warmup can be scheduled using torch_scheduler_kwargs
@@ -373,7 +337,7 @@ metrics:
   #     threshold_kwargs: null
   # l1000_mcf7: *classif_metrics
   # pcba_1328:
-  # # use auroc and averageprecision (non_ipu version) so tha nans are handled correctly
+  # 
   #   - name: auroc
   #     metric: auroc
   #     task: binary
@@ -388,17 +352,17 @@ metrics:
   #     threshold_kwargs: null
   pcqm4m_n4: &pcqm_metrics
     - name: mae
-      metric: mae_ipu
+      metric: mae
       target_nan_mask: null
       multitask_handling: mean-per-label
       threshold_kwargs: null
     - name: pearsonr
-      metric: pearsonr_ipu
+      metric: pearsonr
       threshold_kwargs: null
       target_nan_mask: null
       multitask_handling: mean-per-label
     - name: r2
-      metric: r2_score_ipu
+      metric: r2_score
       threshold_kwargs: null
       target_nan_mask: null
       multitask_handling: mean-per-label
diff --git a/expts/neurips2023_configs/base_config/small.yaml b/expts/neurips2023_configs/base_config/small.yaml
index 4914fdda3..a8b706b2a 100644
--- a/expts/neurips2023_configs/base_config/small.yaml
+++ b/expts/neurips2023_configs/base_config/small.yaml
@@ -6,18 +6,11 @@ constants:
   entity: multitask-gnn
 
 accelerator:
-  type: ipu  # cpu or ipu or gpu
+  type: gpu  # cpu or gpu
   config_override:
     datamodule:
       args:
-        ipu_dataloader_training_opts:
-          mode: async
-          max_num_nodes_per_graph: 44 # train max nodes: 20, max_edges: 54
-          max_num_edges_per_graph: 80
-        ipu_dataloader_inference_opts:
-          mode: async
-          max_num_nodes_per_graph: 44 # valid max nodes: 51, max_edges: 118
-          max_num_edges_per_graph: 80
+
         # Data handling-related
         batch_size_training: 50
         batch_size_inference: 50
@@ -29,26 +22,6 @@ accelerator:
         precision: 16
         accumulate_grad_batches: 4
 
-  ipu_config:
-    - deviceIterations(5) # IPU would require large batches to be ready for the model.
-    - replicationFactor(16)
-    # - enableProfiling("graph_analyser")       # The folder where the profile will be stored
-    # - enableExecutableCaching("pop_compiler_cache")
-    - TensorLocations.numIOTiles(128)
-    - _Popart.set("defaultBufferingDepth", 128)
-    - Precision.enableStochasticRounding(True)
-
-# accelerator:
-#   type: cpu  # cpu or ipu or gpu
-#   config_override:
-#     datamodule:
-#       batch_size_training: 64
-#       batch_size_inference: 256
-#     trainer:
-#       trainer:
-#         precision: 32
-#         accumulate_grad_batches: 1
-
 datamodule:
   module_type: "MultitaskFromSmilesDataModule"
   args: # Matches that in the test_multitask_datamodule.py case.
@@ -247,9 +220,9 @@ predictor:
     tox21: ["auroc"]
     zinc: ["mae"]
   loss_fun:
-    qm9: mae_ipu
-    tox21: bce_ipu
-    zinc: mae_ipu
+    qm9: mae
+    tox21: bce
+    zinc: mae
   random_seed: *seed
   optim_kwargs:
     lr: 4.e-5 # warmup can be scheduled using torch_scheduler_kwargs
@@ -270,28 +243,28 @@ predictor:
 metrics:
   qm9: &qm9_metrics
     - name: mae
-      metric: mae_ipu
+      metric: mae
       target_nan_mask: null
       multitask_handling: flatten
       threshold_kwargs: null
     - name: pearsonr
-      metric: pearsonr_ipu
+      metric: pearsonr
       threshold_kwargs: null
       target_nan_mask: null
       multitask_handling: mean-per-label
     - name: r2_score
-      metric: r2_score_ipu
+      metric: r2_score
       target_nan_mask: null
       multitask_handling: mean-per-label
       threshold_kwargs: null
   tox21:
     - name: auroc
-      metric: auroc_ipu
+      metric: auroc
       task: binary
       multitask_handling: mean-per-label
       threshold_kwargs: null
     - name: avpr
-      metric: average_precision_ipu
+      metric: average_precision
       task: binary
       multitask_handling: mean-per-label
       threshold_kwargs: null
diff --git a/expts/neurips2023_configs/baseline/config_small_gcn_baseline.yaml b/expts/neurips2023_configs/baseline/config_small_gcn_baseline.yaml
index e107fa386..22fbde029 100644
--- a/expts/neurips2023_configs/baseline/config_small_gcn_baseline.yaml
+++ b/expts/neurips2023_configs/baseline/config_small_gcn_baseline.yaml
@@ -1,22 +1,14 @@
-# Testing the gcn model with the PCQMv2 dataset on IPU.
+# Testing the gcn model with the PCQMv2 dataset.
 constants:
   name: &name neurips2023_small_data_gcn
   seed: &seed 3000
   raise_train_error: true   # Whether the code should raise an error if it crashes during training
 
 accelerator:
-  type: ipu  # cpu or ipu or gpu
+  type: gpu  # cpu or gpu
   config_override:
     datamodule:
       args:
-        ipu_dataloader_training_opts:
-          mode: async
-          max_num_nodes_per_graph: 44 # train max nodes: 20, max_edges: 54
-          max_num_edges_per_graph: 80
-        ipu_dataloader_inference_opts:
-          mode: async
-          max_num_nodes_per_graph: 100 # valid max nodes: 51, max_edges: 118
-          max_num_edges_per_graph: 200
         # Data handling-related
         batch_size_training: 5
         batch_size_inference: 2
@@ -28,26 +20,6 @@ accelerator:
         precision: 16
         accumulate_grad_batches: 4
 
-  ipu_config:
-    - deviceIterations(5) # IPU would require large batches to be ready for the model.
-    - replicationFactor(16)
-    # - enableProfiling("graph_analyser")       # The folder where the profile will be stored
-    # - enableExecutableCaching("pop_compiler_cache")
-    - TensorLocations.numIOTiles(128)
-    - _Popart.set("defaultBufferingDepth", 128)
-    - Precision.enableStochasticRounding(True)
-
-# accelerator:
-#   type: cpu  # cpu or ipu or gpu
-#   config_override:
-#     datamodule:
-#       batch_size_training: 64
-#       batch_size_inference: 256
-#     trainer:
-#       trainer:
-#         precision: 32
-#         accumulate_grad_batches: 1
-
 datamodule:
   module_type: "MultitaskFromSmilesDataModule"
   args: # Matches that in the test_multitask_datamodule.py case.
@@ -247,9 +219,9 @@ predictor:
     tox21: ["auroc"]
     zinc: ["mae"]
   loss_fun:
-    qm9: mae_ipu
-    tox21: bce_ipu
-    zinc: mae_ipu
+    qm9: mae
+    tox21: bce
+    zinc: mae
   random_seed: *seed
   optim_kwargs:
     lr: 1.e-4 # warmup can be scheduled using torch_scheduler_kwargs
@@ -270,28 +242,28 @@ predictor:
 metrics:
   qm9: &qm9_metrics
     - name: mae
-      metric: mae_ipu
+      metric: mae
       target_nan_mask: null
       multitask_handling: flatten
       threshold_kwargs: null
     - name: pearsonr
-      metric: pearsonr_ipu
+      metric: pearsonr
       threshold_kwargs: null
       target_nan_mask: null
       multitask_handling: mean-per-label
     - name: r2_score
-      metric: r2_score_ipu
+      metric: r2_score
       target_nan_mask: null
       multitask_handling: mean-per-label
       threshold_kwargs: null
   tox21:
     - name: auroc
-      metric: auroc_ipu
+      metric: auroc
       task: binary
       multitask_handling: mean-per-label
       threshold_kwargs: null
     - name: avpr
-      metric: average_precision_ipu
+      metric: average_precision
       task: binary
       multitask_handling: mean-per-label
       threshold_kwargs: null
diff --git a/expts/neurips2023_configs/baseline/config_small_gin_baseline.yaml b/expts/neurips2023_configs/baseline/config_small_gin_baseline.yaml
index bc96d1057..3dbdca7d7 100644
--- a/expts/neurips2023_configs/baseline/config_small_gin_baseline.yaml
+++ b/expts/neurips2023_configs/baseline/config_small_gin_baseline.yaml
@@ -1,4 +1,4 @@
-# Testing the gin model with the PCQMv2 dataset on IPU.
+# Testing the gin model with the PCQMv2 dataset.
 constants:
   name: &name neurips2023_small_data_gin
   config_override: "expts/neurips2023_configs/baseline/config_small_gcn_baseline.yaml"
diff --git a/expts/neurips2023_configs/baseline/config_small_gine_baseline.yaml b/expts/neurips2023_configs/baseline/config_small_gine_baseline.yaml
index 431235bb4..90389e008 100644
--- a/expts/neurips2023_configs/baseline/config_small_gine_baseline.yaml
+++ b/expts/neurips2023_configs/baseline/config_small_gine_baseline.yaml
@@ -1,4 +1,4 @@
-# Testing the gine model with the PCQMv2 dataset on IPU.
+# Testing the gine model with the PCQMv2 dataset.
 constants:
   name: &name neurips2023_small_data_gine
   config_override: "expts/neurips2023_configs/baseline/config_small_gcn_baseline.yaml"
diff --git a/expts/neurips2023_configs/config_classifigression_l1000.yaml b/expts/neurips2023_configs/config_classifigression_l1000.yaml
index fb77ad457..16d6d1c73 100644
--- a/expts/neurips2023_configs/config_classifigression_l1000.yaml
+++ b/expts/neurips2023_configs/config_classifigression_l1000.yaml
@@ -1,44 +1,11 @@
-# Testing the mpnn only model with the PCQMv2 dataset on IPU.
+# Testing the mpnn only model with the PCQMv2 dataset.
 constants:
   name: &name neurips2023_small_data_mpnn
   seed: &seed 42
   raise_train_error: true   # Whether the code should raise an error if it crashes during training
 
-#accelerator:
-#  type: ipu  # cpu or ipu or gpu
-#  config_override:
-#    datamodule:
-#      args:
-#        ipu_dataloader_training_opts:
-#          mode: async
-#          max_num_nodes_per_graph: 24 # train max nodes: 20, max_edges: 54
-#          max_num_edges_per_graph: 60
-#        ipu_dataloader_inference_opts:
-#          mode: async
-#          max_num_nodes_per_graph: 24 # valid max nodes: 51, max_edges: 118
-#          max_num_edges_per_graph: 60
-#        # Data handling-related
-#        batch_size_training: 50
-#        batch_size_inference: 50
-##    predictor:
-##      optim_kwargs:
-##        loss_scaling: 1024
-#    trainer:
-#      trainer:
-#        precision: 16
-#        accumulate_grad_batches: 4
-#
-#  ipu_config:
-#    - deviceIterations(20) # IPU would require large batches to be ready for the model.
-#    - replicationFactor(16)
-#    # - enableProfiling("graph_analyser")       # The folder where the profile will be stored
-#    # - enableExecutableCaching("pop_compiler_cache")
-#    - TensorLocations.numIOTiles(128)
-#    - _Popart.set("defaultBufferingDepth", 128)
-#    - Precision.enableStochasticRounding(True)
-
 accelerator:
- type: gpu  # cpu or ipu or gpu
+ type: gpu  # cpu or gpu
  config_override:
    datamodule:
      batch_size_training: 64
diff --git a/expts/neurips2023_configs/config_large_gcn.yaml b/expts/neurips2023_configs/config_large_gcn.yaml
index 1219401df..cf56dcc3a 100644
--- a/expts/neurips2023_configs/config_large_gcn.yaml
+++ b/expts/neurips2023_configs/config_large_gcn.yaml
@@ -1,4 +1,4 @@
-# Running the gcn model with the largemix dataset on IPU.
+# Running the gcn model with the largemix dataset.
 
 defaults:
   - base_config: large
diff --git a/expts/neurips2023_configs/config_large_gcn_g25.yaml b/expts/neurips2023_configs/config_large_gcn_g25.yaml
index 35c4e27d5..80a72cc3b 100644
--- a/expts/neurips2023_configs/config_large_gcn_g25.yaml
+++ b/expts/neurips2023_configs/config_large_gcn_g25.yaml
@@ -1,4 +1,4 @@
-# Running the gcn model with the largemix dataset on IPU.
+# Running the gcn model with the largemix dataset.
 
 defaults:
   # - base_config: large
diff --git a/expts/neurips2023_configs/config_large_gcn_gpu.yaml b/expts/neurips2023_configs/config_large_gcn_gpu.yaml
index 31a02e22c..6c5be45fe 100644
--- a/expts/neurips2023_configs/config_large_gcn_gpu.yaml
+++ b/expts/neurips2023_configs/config_large_gcn_gpu.yaml
@@ -59,7 +59,7 @@ architecture:
 
 predictor:
   loss_fun:
-    pcba_1328: bce_logits_ipu
+    pcba_1328: bce_logits
 
   torch_scheduler_kwargs:
     max_num_epochs: &max_epochs 20
diff --git a/expts/neurips2023_configs/config_large_gcn_n4.yaml b/expts/neurips2023_configs/config_large_gcn_n4.yaml
index 61d335c12..616feec09 100644
--- a/expts/neurips2023_configs/config_large_gcn_n4.yaml
+++ b/expts/neurips2023_configs/config_large_gcn_n4.yaml
@@ -1,4 +1,4 @@
-# Running the gcn model with the largemix dataset on IPU.
+# Running the gcn model with the largemix dataset.
 
 defaults:
   # - base_config: large
diff --git a/expts/neurips2023_configs/config_large_gcn_pcba.yaml b/expts/neurips2023_configs/config_large_gcn_pcba.yaml
index f11d8595d..d95401cf6 100644
--- a/expts/neurips2023_configs/config_large_gcn_pcba.yaml
+++ b/expts/neurips2023_configs/config_large_gcn_pcba.yaml
@@ -1,4 +1,4 @@
-# Running the gcn model with the largemix dataset on IPU.
+# Running the gcn model with the largemix dataset.
 
 defaults:
   # - base_config: large
diff --git a/expts/neurips2023_configs/config_large_gin.yaml b/expts/neurips2023_configs/config_large_gin.yaml
index 6c3f47898..b2a0186f3 100644
--- a/expts/neurips2023_configs/config_large_gin.yaml
+++ b/expts/neurips2023_configs/config_large_gin.yaml
@@ -1,4 +1,4 @@
-# Running the gin model with the largemix dataset on IPU.
+# Running the gin model with the largemix dataset.
 defaults:
   - base_config: large
   - _self_
diff --git a/expts/neurips2023_configs/config_large_gin_g25.yaml b/expts/neurips2023_configs/config_large_gin_g25.yaml
index 538e994b1..702e9fe68 100644
--- a/expts/neurips2023_configs/config_large_gin_g25.yaml
+++ b/expts/neurips2023_configs/config_large_gin_g25.yaml
@@ -1,4 +1,4 @@
-# Running the gin model with the largemix dataset on IPU.
+# Running the gin model with the largemix dataset.
 defaults:
   # - base_config: large
   - base_config: large_pcqm_g25
diff --git a/expts/neurips2023_configs/config_large_gin_n4.yaml b/expts/neurips2023_configs/config_large_gin_n4.yaml
index c51e0e07d..3e41cf95b 100644
--- a/expts/neurips2023_configs/config_large_gin_n4.yaml
+++ b/expts/neurips2023_configs/config_large_gin_n4.yaml
@@ -1,4 +1,4 @@
-# Running the gin model with the largemix dataset on IPU.
+# Running the gin model with the largemix dataset.
 defaults:
   # - base_config: large
   # - base_config: large_pcqm_g25
diff --git a/expts/neurips2023_configs/config_large_gin_pcba.yaml b/expts/neurips2023_configs/config_large_gin_pcba.yaml
index 8bd33609b..af0e4945e 100644
--- a/expts/neurips2023_configs/config_large_gin_pcba.yaml
+++ b/expts/neurips2023_configs/config_large_gin_pcba.yaml
@@ -1,4 +1,4 @@
-# Running the gin model with the largemix dataset on IPU.
+# Running the gin model with the largemix dataset.
 defaults:
   # - base_config: large
   # - base_config: large_pcqm_g25
diff --git a/expts/neurips2023_configs/config_large_gine.yaml b/expts/neurips2023_configs/config_large_gine.yaml
index 793304ce0..6f82d3233 100644
--- a/expts/neurips2023_configs/config_large_gine.yaml
+++ b/expts/neurips2023_configs/config_large_gine.yaml
@@ -1,4 +1,4 @@
-# Running the gine model with the largemix dataset on IPU.
+# Running the gine model with the largemix dataset.
 
 defaults:
   - base_config: large
diff --git a/expts/neurips2023_configs/config_large_gine_g25.yaml b/expts/neurips2023_configs/config_large_gine_g25.yaml
index e8002be3b..cceaa448f 100644
--- a/expts/neurips2023_configs/config_large_gine_g25.yaml
+++ b/expts/neurips2023_configs/config_large_gine_g25.yaml
@@ -1,4 +1,4 @@
-# Running the gine model with the largemix dataset on IPU.
+# Running the gine model with the largemix dataset.
 
 defaults:
   # - base_config: large
diff --git a/expts/neurips2023_configs/config_large_gine_n4.yaml b/expts/neurips2023_configs/config_large_gine_n4.yaml
index df07380a4..d298fc183 100644
--- a/expts/neurips2023_configs/config_large_gine_n4.yaml
+++ b/expts/neurips2023_configs/config_large_gine_n4.yaml
@@ -1,4 +1,4 @@
-# Running the gine model with the largemix dataset on IPU.
+# Running the gine model with the largemix dataset.
 
 defaults:
   # - base_config: large
diff --git a/expts/neurips2023_configs/config_large_gine_pcba.yaml b/expts/neurips2023_configs/config_large_gine_pcba.yaml
index 505935a57..668b7dfc2 100644
--- a/expts/neurips2023_configs/config_large_gine_pcba.yaml
+++ b/expts/neurips2023_configs/config_large_gine_pcba.yaml
@@ -1,4 +1,4 @@
-# Running the gine model with the largemix dataset on IPU.
+# Running the gine model with the largemix dataset.
 
 defaults:
   # - base_config: large
diff --git a/expts/neurips2023_configs/config_large_mpnn.yaml b/expts/neurips2023_configs/config_large_mpnn.yaml
index 365927473..ca280f68a 100644
--- a/expts/neurips2023_configs/config_large_mpnn.yaml
+++ b/expts/neurips2023_configs/config_large_mpnn.yaml
@@ -1,4 +1,4 @@
-# Running the mpnn model with the largemix dataset on IPU.
+# Running the mpnn model with the largemix dataset.
 
 defaults:
  - base_config: large
diff --git a/expts/neurips2023_configs/config_luis_jama.yaml b/expts/neurips2023_configs/config_luis_jama.yaml
index e0549e0f0..dc25ed212 100644
--- a/expts/neurips2023_configs/config_luis_jama.yaml
+++ b/expts/neurips2023_configs/config_luis_jama.yaml
@@ -1,44 +1,11 @@
-# Testing the mpnn only model with the PCQMv2 dataset on IPU.
+# Testing the mpnn only model with the PCQMv2 dataset.
 constants:
   name: &name neurips2023_small_data_mpnn
   seed: &seed 42
   raise_train_error: true   # Whether the code should raise an error if it crashes during training
 
-# accelerator:
-#   type: ipu  # cpu or ipu or gpu
-#   config_override:
-#     datamodule:
-#       args:
-#         ipu_dataloader_training_opts:
-#           mode: async
-#           max_num_nodes_per_graph: 24 # train max nodes: 20, max_edges: 54
-#           max_num_edges_per_graph: 60
-#         ipu_dataloader_inference_opts:
-#           mode: async
-#           max_num_nodes_per_graph: 24 # valid max nodes: 51, max_edges: 118
-#           max_num_edges_per_graph: 60
-#         # Data handling-related
-#         batch_size_training: 50
-#         batch_size_inference: 50
-#     predictor:
-#       optim_kwargs:
-#         loss_scaling: 1024
-#     trainer:
-#       trainer:
-#         precision: 16
-#         accumulate_grad_batches: 4
-
-  # ipu_config:
-  #   - deviceIterations(20) # IPU would require large batches to be ready for the model.
-  #   - replicationFactor(16)
-  #   # - enableProfiling("graph_analyser")       # The folder where the profile will be stored
-  #   # - enableExecutableCaching("pop_compiler_cache")
-  #   - TensorLocations.numIOTiles(128)
-  #   - _Popart.set("defaultBufferingDepth", 128)
-  #   - Precision.enableStochasticRounding(True)
-
 accelerator:
-  type: cpu  # cpu or ipu or gpu
+  type: cpu  # cpu or gpu
   config_override:
     datamodule:
       batch_size_training: 64
@@ -251,8 +218,8 @@ predictor:
     pcqm20k_g13: []
     pcqm20k_n4: []
   loss_fun:
-    pcqm20k_g13: mae_ipu
-    pcqm20k_n4: mae_ipu
+    pcqm20k_g13: mae
+    pcqm20k_n4: mae
   random_seed: *seed
   optim_kwargs:
     lr: 4.e-5 # warmup can be scheduled using torch_scheduler_kwargs
@@ -273,12 +240,12 @@ predictor:
 metrics:
   pcqm20k_g13: &pcqm_metrics
     - name: mae
-      metric: mae_ipu
+      metric: mae
       target_nan_mask: null
       multitask_handling: flatten
       threshold_kwargs: null
     - name: pearsonr
-      metric: pearsonr_ipu
+      metric: pearsonr
       threshold_kwargs: null
       target_nan_mask: null
       multitask_handling: mean-per-label
diff --git a/expts/neurips2023_configs/config_small_gated_gcn.yaml b/expts/neurips2023_configs/config_small_gated_gcn.yaml
index 8e00d26f6..5a542d96d 100644
--- a/expts/neurips2023_configs/config_small_gated_gcn.yaml
+++ b/expts/neurips2023_configs/config_small_gated_gcn.yaml
@@ -1,4 +1,4 @@
-# Testing the gated_gcn model with the PCQMv2 dataset on IPU.
+# Testing the gated_gcn model with the PCQMv2 dataset.
 
 defaults:
   - base_config: small
diff --git a/expts/neurips2023_configs/config_small_gcn.yaml b/expts/neurips2023_configs/config_small_gcn.yaml
index 114ce26dc..d43080a4c 100644
--- a/expts/neurips2023_configs/config_small_gcn.yaml
+++ b/expts/neurips2023_configs/config_small_gcn.yaml
@@ -1,4 +1,4 @@
-# Testing the gcn model with the toymix dataset on IPU.
+# Testing the gcn model with the toymix dataset.
 
 defaults:
   - base_config: small
diff --git a/expts/neurips2023_configs/config_small_gcn_gpu.yaml b/expts/neurips2023_configs/config_small_gcn_gpu.yaml
index ccad70af6..e1223da2f 100644
--- a/expts/neurips2023_configs/config_small_gcn_gpu.yaml
+++ b/expts/neurips2023_configs/config_small_gcn_gpu.yaml
@@ -12,7 +12,7 @@ architecture:
     layer_type: 'pyg:gcn' #pyg:gine #'pyg:gps' # pyg:gated-gcn, pyg:gine,pyg:gps
 
 accelerator:
-  type: gpu  # cpu or ipu or gpu
+  type: gpu  # cpu or gpu
   float32_matmul_precision: medium
   config_override:
     datamodule:
@@ -51,7 +51,7 @@ architecture:
 
 predictor:
   loss_fun:
-    tox21: bce_logits_ipu
+    tox21: bce_logits
 
   torch_scheduler_kwargs:
     max_num_epochs: &max_epochs 300
diff --git a/expts/neurips2023_configs/config_small_gin.yaml b/expts/neurips2023_configs/config_small_gin.yaml
index e018f722a..ff86251f7 100644
--- a/expts/neurips2023_configs/config_small_gin.yaml
+++ b/expts/neurips2023_configs/config_small_gin.yaml
@@ -1,4 +1,4 @@
-# Testing the gin model with the PCQMv2 dataset on IPU.
+# Testing the gin model with the PCQMv2 dataset.
 
 defaults:
   - base_config: small
diff --git a/expts/neurips2023_configs/config_small_gine.yaml b/expts/neurips2023_configs/config_small_gine.yaml
index 111bebbc2..4ec66a4f2 100644
--- a/expts/neurips2023_configs/config_small_gine.yaml
+++ b/expts/neurips2023_configs/config_small_gine.yaml
@@ -1,4 +1,4 @@
-# Testing the gine model with the PCQMv2 dataset on IPU.
+# Testing the gine model with the PCQMv2 dataset.
 
 defaults:
   - base_config: small
diff --git a/expts/neurips2023_configs/config_small_mpnn.yaml b/expts/neurips2023_configs/config_small_mpnn.yaml
index 357a8f560..12c7a17e1 100644
--- a/expts/neurips2023_configs/config_small_mpnn.yaml
+++ b/expts/neurips2023_configs/config_small_mpnn.yaml
@@ -1,4 +1,4 @@
-# Testing the mpnn only model with the PCQMv2 dataset on IPU.
+# Testing the mpnn only model with the PCQMv2 dataset.
 
 defaults:
  - base_config: small
diff --git a/expts/neurips2023_configs/single_task_gcn/config_large_gcn_mcf7.yaml b/expts/neurips2023_configs/single_task_gcn/config_large_gcn_mcf7.yaml
index 1aba37eb4..d4dc601fa 100644
--- a/expts/neurips2023_configs/single_task_gcn/config_large_gcn_mcf7.yaml
+++ b/expts/neurips2023_configs/single_task_gcn/config_large_gcn_mcf7.yaml
@@ -5,18 +5,10 @@ constants:
   raise_train_error: true   # Whether the code should raise an error if it crashes during training
 
 accelerator:
-  type: ipu  # cpu or ipu or gpu
+  type: gpu  # cpu or gpu
   config_override:
     datamodule:
       args:
-        ipu_dataloader_training_opts:
-          mode: async
-          max_num_nodes_per_graph: 60 # train max nodes: 20, max_edges: 54
-          max_num_edges_per_graph: 100
-        ipu_dataloader_inference_opts:
-          mode: async
-          max_num_nodes_per_graph: 60 # valid max nodes: 51, max_edges: 118
-          max_num_edges_per_graph: 100
         # Data handling-related
         batch_size_training: 10
         batch_size_inference: 2
@@ -28,26 +20,6 @@ accelerator:
         precision: 32
         accumulate_grad_batches: 8
 
-  ipu_config:
-    - deviceIterations(1) # IPU would require large batches to be ready for the model.
-    - replicationFactor(16)
-    # - enableProfiling("graph_analyser")       # The folder where the profile will be stored
-    # - enableExecutableCaching("pop_compiler_cache")
-    - TensorLocations.numIOTiles(128)
-    - _Popart.set("defaultBufferingDepth", 128)
-    - Precision.enableStochasticRounding(True)
-
-# accelerator:
-#   type: cpu  # cpu or ipu or gpu
-#   config_override:
-#     datamodule:
-#       batch_size_training: 64
-#       batch_size_inference: 256
-#     trainer:
-#       trainer:
-#         precision: 32
-#         accumulate_grad_batches: 1
-
 datamodule:
   module_type: "MultitaskFromSmilesDataModule"
   args: # Matches that in the test_multitask_datamodule.py case.
@@ -196,7 +168,7 @@ predictor:
     l1000_mcf7: []
   loss_fun:
     l1000_mcf7:
-      name: hybrid_ce_ipu
+      name: hybrid_ce
       n_brackets: 5
   random_seed: *seed
   optim_kwargs:
@@ -218,13 +190,13 @@ predictor:
 metrics:
   l1000_mcf7: &classif_metrics
     - name: auroc
-      metric: auroc_ipu
+      metric: auroc
       num_classes: 5
       task: multiclass
       multitask_handling: mean-per-label
       threshold_kwargs: null
     - name: avpr
-      metric: average_precision_ipu
+      metric: average_precision
       num_classes: 5
       task: multiclass
       target_to_int: True
diff --git a/expts/neurips2023_configs/single_task_gcn/config_large_gcn_pcba.yaml b/expts/neurips2023_configs/single_task_gcn/config_large_gcn_pcba.yaml
index 77837d750..a59215b90 100644
--- a/expts/neurips2023_configs/single_task_gcn/config_large_gcn_pcba.yaml
+++ b/expts/neurips2023_configs/single_task_gcn/config_large_gcn_pcba.yaml
@@ -5,18 +5,10 @@ constants:
   raise_train_error: true   # Whether the code should raise an error if it crashes during training
 
 accelerator:
-  type: ipu  # cpu or ipu or gpu
+  type: gpu  # cpu or gpu
   config_override:
     datamodule:
       args:
-        ipu_dataloader_training_opts:
-          mode: async
-          max_num_nodes_per_graph: 60 # train max nodes: 20, max_edges: 54
-          max_num_edges_per_graph: 100
-        ipu_dataloader_inference_opts:
-          mode: async
-          max_num_nodes_per_graph: 200 # valid max nodes: 51, max_edges: 118
-          max_num_edges_per_graph: 400
         # Data handling-related
         batch_size_training: 10
         batch_size_inference: 2
@@ -28,26 +20,6 @@ accelerator:
         precision: 32
         accumulate_grad_batches: 8
 
-  ipu_config:
-    - deviceIterations(10) # IPU would require large batches to be ready for the model.
-    - replicationFactor(16)
-    # - enableProfiling("graph_analyser")       # The folder where the profile will be stored
-    # - enableExecutableCaching("pop_compiler_cache")
-    - TensorLocations.numIOTiles(128)
-    - _Popart.set("defaultBufferingDepth", 128)
-    - Precision.enableStochasticRounding(True)
-
-# accelerator:
-#   type: cpu  # cpu or ipu or gpu
-#   config_override:
-#     datamodule:
-#       batch_size_training: 64
-#       batch_size_inference: 256
-#     trainer:
-#       trainer:
-#         precision: 32
-#         accumulate_grad_batches: 1
-
 datamodule:
   module_type: "MultitaskFromSmilesDataModule"
   args: # Matches that in the test_multitask_datamodule.py case.
@@ -193,7 +165,7 @@ predictor:
   metrics_on_training_set:
     pcba_1328: []
   loss_fun:
-    pcba_1328: bce_ipu
+    pcba_1328: bce
   random_seed: *seed
   optim_kwargs:
     lr: 1.e-4 # warmup can be scheduled using torch_scheduler_kwargs
@@ -214,12 +186,12 @@ predictor:
 metrics:
   pcba_1328:
     - name: auroc
-      metric: auroc_ipu
+      metric: auroc
       task: binary
       multitask_handling: mean-per-label
       threshold_kwargs: null
     - name: avpr
-      metric: average_precision_ipu
+      metric: average_precision
       task: binary
       multitask_handling: mean-per-label
       threshold_kwargs: null
diff --git a/expts/neurips2023_configs/single_task_gcn/config_large_gcn_vcap.yaml b/expts/neurips2023_configs/single_task_gcn/config_large_gcn_vcap.yaml
index 1c021a559..4a260e3ee 100644
--- a/expts/neurips2023_configs/single_task_gcn/config_large_gcn_vcap.yaml
+++ b/expts/neurips2023_configs/single_task_gcn/config_large_gcn_vcap.yaml
@@ -5,18 +5,10 @@ constants:
   raise_train_error: true   # Whether the code should raise an error if it crashes during training
 
 accelerator:
-  type: ipu  # cpu or ipu or gpu
+  type: gpu  # cpu or gpu
   config_override:
     datamodule:
       args:
-        ipu_dataloader_training_opts:
-          mode: async
-          max_num_nodes_per_graph: 60 # train max nodes: 20, max_edges: 54
-          max_num_edges_per_graph: 100
-        ipu_dataloader_inference_opts:
-          mode: async
-          max_num_nodes_per_graph: 60 # valid max nodes: 51, max_edges: 118
-          max_num_edges_per_graph: 150
         # Data handling-related
         batch_size_training: 10
         batch_size_inference: 2
@@ -28,26 +20,6 @@ accelerator:
         precision: 32
         accumulate_grad_batches: 8
 
-  ipu_config:
-    - deviceIterations(1) # IPU would require large batches to be ready for the model.
-    - replicationFactor(16)
-    # - enableProfiling("graph_analyser")       # The folder where the profile will be stored
-    # - enableExecutableCaching("pop_compiler_cache")
-    - TensorLocations.numIOTiles(128)
-    - _Popart.set("defaultBufferingDepth", 128)
-    - Precision.enableStochasticRounding(True)
-
-# accelerator:
-#   type: cpu  # cpu or ipu or gpu
-#   config_override:
-#     datamodule:
-#       batch_size_training: 64
-#       batch_size_inference: 256
-#     trainer:
-#       trainer:
-#         precision: 32
-#         accumulate_grad_batches: 1
-
 datamodule:
   module_type: "MultitaskFromSmilesDataModule"
   args: # Matches that in the test_multitask_datamodule.py case.
@@ -194,7 +166,7 @@ predictor:
     l1000_vcap: []
   loss_fun:
     l1000_vcap:
-      name: hybrid_ce_ipu
+      name: hybrid_ce
       n_brackets: 5
   random_seed: *seed
   optim_kwargs:
@@ -216,13 +188,13 @@ predictor:
 metrics:
   l1000_vcap: &classif_metrics
     - name: auroc
-      metric: auroc_ipu
+      metric: auroc
       num_classes: 5
       task: multiclass
       multitask_handling: mean-per-label
       threshold_kwargs: null
     - name: avpr
-      metric: average_precision_ipu
+      metric: average_precision
       num_classes: 5
       task: multiclass
       target_to_int: True
diff --git a/expts/neurips2023_configs/single_task_gin/config_large_gin_g25.yaml b/expts/neurips2023_configs/single_task_gin/config_large_gin_g25.yaml
index bd09385f5..2e4cdaf53 100644
--- a/expts/neurips2023_configs/single_task_gin/config_large_gin_g25.yaml
+++ b/expts/neurips2023_configs/single_task_gin/config_large_gin_g25.yaml
@@ -5,18 +5,10 @@ constants:
   raise_train_error: true   # Whether the code should raise an error if it crashes during training
 
 accelerator:
-  type: ipu  # cpu or ipu or gpu
+  type: gpu  # cpu or gpu
   config_override:
     datamodule:
       args:
-        ipu_dataloader_training_opts:
-          mode: async
-          max_num_nodes_per_graph: 30 # train max nodes: 20, max_edges: 54
-          max_num_edges_per_graph: 100
-        ipu_dataloader_inference_opts:
-          mode: async
-          max_num_nodes_per_graph: 30 # valid max nodes: 51, max_edges: 118
-          max_num_edges_per_graph: 100
         # Data handling-related
         batch_size_training: 10
         batch_size_inference: 10
@@ -28,26 +20,6 @@ accelerator:
         precision: 32
         accumulate_grad_batches: 8
 
-  ipu_config:
-    - deviceIterations(10) # IPU would require large batches to be ready for the model.
-    - replicationFactor(16)
-    # - enableProfiling("graph_analyser")       # The folder where the profile will be stored
-    # - enableExecutableCaching("pop_compiler_cache")
-    - TensorLocations.numIOTiles(128)
-    - _Popart.set("defaultBufferingDepth", 128)
-    - Precision.enableStochasticRounding(True)
-
-# accelerator:
-#   type: cpu  # cpu or ipu or gpu
-#   config_override:
-#     datamodule:
-#       batch_size_training: 64
-#       batch_size_inference: 256
-#     trainer:
-#       trainer:
-#         precision: 32
-#         accumulate_grad_batches: 1
-
 datamodule:
   module_type: "MultitaskFromSmilesDataModule"
   args: # Matches that in the test_multitask_datamodule.py case.
@@ -196,7 +168,7 @@ predictor:
   metrics_on_training_set:
     pcqm4m_g25: []
   loss_fun:
-    pcqm4m_g25: mae_ipu
+    pcqm4m_g25: mae
   random_seed: *seed
   optim_kwargs:
     lr: 1.e-4 # warmup can be scheduled using torch_scheduler_kwargs
@@ -217,17 +189,17 @@ predictor:
 metrics:
   pcqm4m_g25: &pcqm_metrics
     - name: mae
-      metric: mae_ipu
+      metric: mae
       target_nan_mask: null
       multitask_handling: flatten
       threshold_kwargs: null
     - name: pearsonr
-      metric: pearsonr_ipu
+      metric: pearsonr
       threshold_kwargs: null
       target_nan_mask: null
       multitask_handling: mean-per-label
     - name: r2
-      metric: r2_score_ipu
+      metric: r2_score
       threshold_kwargs: null
       target_nan_mask: null
       multitask_handling: mean-per-label
diff --git a/expts/neurips2023_configs/single_task_gin/config_large_gin_mcf7.yaml b/expts/neurips2023_configs/single_task_gin/config_large_gin_mcf7.yaml
index 5abf9790d..09425ec37 100644
--- a/expts/neurips2023_configs/single_task_gin/config_large_gin_mcf7.yaml
+++ b/expts/neurips2023_configs/single_task_gin/config_large_gin_mcf7.yaml
@@ -1,22 +1,14 @@
-# Testing the gine model with the PCQMv2 dataset on IPU.
+# Testing the gine model with the PCQMv2 dataset.
 constants:
   name: &name neurips2023_large_data_gine_mcf7
   seed: &seed 42
   raise_train_error: true   # Whether the code should raise an error if it crashes during training
 
 accelerator:
-  type: ipu  # cpu or ipu or gpu
+  type: gpu  # cpu or gpu
   config_override:
     datamodule:
       args:
-        ipu_dataloader_training_opts:
-          mode: async
-          max_num_nodes_per_graph: 60 # train max nodes: 20, max_edges: 54
-          max_num_edges_per_graph: 100
-        ipu_dataloader_inference_opts:
-          mode: async
-          max_num_nodes_per_graph: 60 # valid max nodes: 51, max_edges: 118
-          max_num_edges_per_graph: 100
         # Data handling-related
         batch_size_training: 10
         batch_size_inference: 2
@@ -28,26 +20,6 @@ accelerator:
         precision: 32
         accumulate_grad_batches: 8
 
-  ipu_config:
-    - deviceIterations(1) # IPU would require large batches to be ready for the model.
-    - replicationFactor(16)
-    # - enableProfiling("graph_analyser")       # The folder where the profile will be stored
-    # - enableExecutableCaching("pop_compiler_cache")
-    - TensorLocations.numIOTiles(128)
-    - _Popart.set("defaultBufferingDepth", 128)
-    - Precision.enableStochasticRounding(True)
-
-# accelerator:
-#   type: cpu  # cpu or ipu or gpu
-#   config_override:
-#     datamodule:
-#       batch_size_training: 64
-#       batch_size_inference: 256
-#     trainer:
-#       trainer:
-#         precision: 32
-#         accumulate_grad_batches: 1
-
 datamodule:
   module_type: "MultitaskFromSmilesDataModule"
   args: # Matches that in the test_multitask_datamodule.py case.
@@ -196,7 +168,7 @@ predictor:
     l1000_mcf7: []
   loss_fun:
     l1000_mcf7:
-      name: hybrid_ce_ipu
+      name: hybrid_ce
       n_brackets: 5
   random_seed: *seed
   optim_kwargs:
@@ -218,13 +190,13 @@ predictor:
 metrics:
   l1000_mcf7: &classif_metrics
     - name: auroc
-      metric: auroc_ipu
+      metric: auroc
       num_classes: 5
       task: multiclass
       multitask_handling: mean-per-label
       threshold_kwargs: null
     - name: avpr
-      metric: average_precision_ipu
+      metric: average_precision
       num_classes: 5
       task: multiclass
       target_to_int: True
diff --git a/expts/neurips2023_configs/single_task_gin/config_large_gin_n4.yaml b/expts/neurips2023_configs/single_task_gin/config_large_gin_n4.yaml
index 834967498..9ef7254fc 100644
--- a/expts/neurips2023_configs/single_task_gin/config_large_gin_n4.yaml
+++ b/expts/neurips2023_configs/single_task_gin/config_large_gin_n4.yaml
@@ -5,18 +5,10 @@ constants:
   raise_train_error: true   # Whether the code should raise an error if it crashes during training
 
 accelerator:
-  type: ipu  # cpu or ipu or gpu
+  type: gpu  # cpu or gpu
   config_override:
     datamodule:
       args:
-        ipu_dataloader_training_opts:
-          mode: async
-          max_num_nodes_per_graph: 30 # train max nodes: 20, max_edges: 54
-          max_num_edges_per_graph: 100
-        ipu_dataloader_inference_opts:
-          mode: async
-          max_num_nodes_per_graph: 30 # valid max nodes: 51, max_edges: 118
-          max_num_edges_per_graph: 100
         # Data handling-related
         batch_size_training: 10
         batch_size_inference: 10
@@ -28,26 +20,6 @@ accelerator:
         precision: 32
         accumulate_grad_batches: 8
 
-  ipu_config:
-    - deviceIterations(10) # IPU would require large batches to be ready for the model.
-    - replicationFactor(16)
-    # - enableProfiling("graph_analyser")       # The folder where the profile will be stored
-    # - enableExecutableCaching("pop_compiler_cache")
-    - TensorLocations.numIOTiles(128)
-    - _Popart.set("defaultBufferingDepth", 128)
-    - Precision.enableStochasticRounding(True)
-
-# accelerator:
-#   type: cpu  # cpu or ipu or gpu
-#   config_override:
-#     datamodule:
-#       batch_size_training: 64
-#       batch_size_inference: 256
-#     trainer:
-#       trainer:
-#         precision: 32
-#         accumulate_grad_batches: 1
-
 datamodule:
   module_type: "MultitaskFromSmilesDataModule"
   args: # Matches that in the test_multitask_datamodule.py case.
@@ -197,7 +169,7 @@ predictor:
   metrics_on_training_set:
     pcqm4m_n4: []
   loss_fun:
-    pcqm4m_n4: mae_ipu
+    pcqm4m_n4: mae
   random_seed: *seed
   optim_kwargs:
     lr: 1.e-4 # warmup can be scheduled using torch_scheduler_kwargs
@@ -218,17 +190,17 @@ predictor:
 metrics:
   pcqm4m_n4: &pcqm_metrics
     - name: mae
-      metric: mae_ipu
+      metric: mae
       target_nan_mask: null
       multitask_handling: flatten
       threshold_kwargs: null
     - name: pearsonr
-      metric: pearsonr_ipu
+      metric: pearsonr
       threshold_kwargs: null
       target_nan_mask: null
       multitask_handling: mean-per-label
     - name: r2
-      metric: r2_score_ipu
+      metric: r2_score
       threshold_kwargs: null
       target_nan_mask: null
       multitask_handling: mean-per-label
diff --git a/expts/neurips2023_configs/single_task_gin/config_large_gin_pcba.yaml b/expts/neurips2023_configs/single_task_gin/config_large_gin_pcba.yaml
index f390a7a2b..ecf18ce9b 100644
--- a/expts/neurips2023_configs/single_task_gin/config_large_gin_pcba.yaml
+++ b/expts/neurips2023_configs/single_task_gin/config_large_gin_pcba.yaml
@@ -5,18 +5,9 @@ constants:
   raise_train_error: true   # Whether the code should raise an error if it crashes during training
 
 accelerator:
-  type: ipu  # cpu or ipu or gpu
+  type: gpu  # cpu or gpu
   config_override:
     datamodule:
-      args:
-        ipu_dataloader_training_opts:
-          mode: async
-          max_num_nodes_per_graph: 60 # train max nodes: 20, max_edges: 54
-          max_num_edges_per_graph: 100
-        ipu_dataloader_inference_opts:
-          mode: async
-          max_num_nodes_per_graph: 200 # valid max nodes: 51, max_edges: 118
-          max_num_edges_per_graph: 400
         # Data handling-related
         batch_size_training: 10
         batch_size_inference: 2
@@ -28,26 +19,6 @@ accelerator:
         precision: 32
         accumulate_grad_batches: 8
 
-  ipu_config:
-    - deviceIterations(10) # IPU would require large batches to be ready for the model.
-    - replicationFactor(16)
-    # - enableProfiling("graph_analyser")       # The folder where the profile will be stored
-    # - enableExecutableCaching("pop_compiler_cache")
-    - TensorLocations.numIOTiles(128)
-    - _Popart.set("defaultBufferingDepth", 128)
-    - Precision.enableStochasticRounding(True)
-
-# accelerator:
-#   type: cpu  # cpu or ipu or gpu
-#   config_override:
-#     datamodule:
-#       batch_size_training: 64
-#       batch_size_inference: 256
-#     trainer:
-#       trainer:
-#         precision: 32
-#         accumulate_grad_batches: 1
-
 datamodule:
   module_type: "MultitaskFromSmilesDataModule"
   args: # Matches that in the test_multitask_datamodule.py case.
@@ -193,7 +164,7 @@ predictor:
   metrics_on_training_set:
     pcba_1328: []
   loss_fun:
-    pcba_1328: bce_ipu
+    pcba_1328: bce
   random_seed: *seed
   optim_kwargs:
     lr: 1.e-4 # warmup can be scheduled using torch_scheduler_kwargs
@@ -214,12 +185,12 @@ predictor:
 metrics:
   pcba_1328:
     - name: auroc
-      metric: auroc_ipu
+      metric: auroc
       task: binary
       multitask_handling: mean-per-label
       threshold_kwargs: null
     - name: avpr
-      metric: average_precision_ipu
+      metric: average_precision
       task: binary
       multitask_handling: mean-per-label
       threshold_kwargs: null
diff --git a/expts/neurips2023_configs/single_task_gin/config_large_gin_pcq.yaml b/expts/neurips2023_configs/single_task_gin/config_large_gin_pcq.yaml
index d13a757f3..13a44e1c0 100644
--- a/expts/neurips2023_configs/single_task_gin/config_large_gin_pcq.yaml
+++ b/expts/neurips2023_configs/single_task_gin/config_large_gin_pcq.yaml
@@ -5,18 +5,10 @@ constants:
   raise_train_error: true   # Whether the code should raise an error if it crashes during training
 
 accelerator:
-  type: ipu  # cpu or ipu or gpu
+  type: gpu  # cpu or gpu
   config_override:
     datamodule:
       args:
-        ipu_dataloader_training_opts:
-          mode: async
-          max_num_nodes_per_graph: 30 # train max nodes: 20, max_edges: 54
-          max_num_edges_per_graph: 100
-        ipu_dataloader_inference_opts:
-          mode: async
-          max_num_nodes_per_graph: 30 # valid max nodes: 51, max_edges: 118
-          max_num_edges_per_graph: 100
         # Data handling-related
         batch_size_training: 10
         batch_size_inference: 10
@@ -28,26 +20,6 @@ accelerator:
         precision: 32
         accumulate_grad_batches: 8
 
-  ipu_config:
-    - deviceIterations(10) # IPU would require large batches to be ready for the model.
-    - replicationFactor(16)
-    # - enableProfiling("graph_analyser")       # The folder where the profile will be stored
-    # - enableExecutableCaching("pop_compiler_cache")
-    - TensorLocations.numIOTiles(128)
-    - _Popart.set("defaultBufferingDepth", 128)
-    - Precision.enableStochasticRounding(True)
-
-# accelerator:
-#   type: cpu  # cpu or ipu or gpu
-#   config_override:
-#     datamodule:
-#       batch_size_training: 64
-#       batch_size_inference: 256
-#     trainer:
-#       trainer:
-#         precision: 32
-#         accumulate_grad_batches: 1
-
 datamodule:
   module_type: "MultitaskFromSmilesDataModule"
   args: # Matches that in the test_multitask_datamodule.py case.
@@ -235,8 +207,8 @@ predictor:
     pcqm4m_g25: []
     pcqm4m_n4: []
   loss_fun:
-    pcqm4m_g25: mae_ipu
-    pcqm4m_n4: mae_ipu
+    pcqm4m_g25: mae
+    pcqm4m_n4: mae
   random_seed: *seed
   optim_kwargs:
     lr: 1.e-4 # warmup can be scheduled using torch_scheduler_kwargs
@@ -257,17 +229,17 @@ predictor:
 metrics:
   pcqm4m_g25: &pcqm_metrics
     - name: mae
-      metric: mae_ipu
+      metric: mae
       target_nan_mask: null
       multitask_handling: flatten
       threshold_kwargs: null
     - name: pearsonr
-      metric: pearsonr_ipu
+      metric: pearsonr
       threshold_kwargs: null
       target_nan_mask: null
       multitask_handling: mean-per-label
     - name: r2
-      metric: r2_score_ipu
+      metric: r2_score
       threshold_kwargs: null
       target_nan_mask: null
       multitask_handling: mean-per-label
diff --git a/expts/neurips2023_configs/single_task_gin/config_large_gin_vcap.yaml b/expts/neurips2023_configs/single_task_gin/config_large_gin_vcap.yaml
index 75f802926..3e716d5b1 100644
--- a/expts/neurips2023_configs/single_task_gin/config_large_gin_vcap.yaml
+++ b/expts/neurips2023_configs/single_task_gin/config_large_gin_vcap.yaml
@@ -1,22 +1,14 @@
-# Testing the gine model with the PCQMv2 dataset on IPU.
+# Testing the gine model with the PCQMv2 dataset.
 constants:
   name: &name neurips2023_large_data_gine_vcap
   seed: &seed 42
   raise_train_error: true   # Whether the code should raise an error if it crashes during training
 
 accelerator:
-  type: ipu  # cpu or ipu or gpu
+  type: gpu  # cpu or gpu
   config_override:
     datamodule:
       args:
-        ipu_dataloader_training_opts:
-          mode: async
-          max_num_nodes_per_graph: 60 # train max nodes: 20, max_edges: 54
-          max_num_edges_per_graph: 100
-        ipu_dataloader_inference_opts:
-          mode: async
-          max_num_nodes_per_graph: 60 # valid max nodes: 51, max_edges: 118
-          max_num_edges_per_graph: 150
         # Data handling-related
         batch_size_training: 10
         batch_size_inference: 2
@@ -28,26 +20,6 @@ accelerator:
         precision: 32
         accumulate_grad_batches: 8
 
-  ipu_config:
-    - deviceIterations(1) # IPU would require large batches to be ready for the model.
-    - replicationFactor(16)
-    # - enableProfiling("graph_analyser")       # The folder where the profile will be stored
-    # - enableExecutableCaching("pop_compiler_cache")
-    - TensorLocations.numIOTiles(128)
-    - _Popart.set("defaultBufferingDepth", 128)
-    - Precision.enableStochasticRounding(True)
-
-# accelerator:
-#   type: cpu  # cpu or ipu or gpu
-#   config_override:
-#     datamodule:
-#       batch_size_training: 64
-#       batch_size_inference: 256
-#     trainer:
-#       trainer:
-#         precision: 32
-#         accumulate_grad_batches: 1
-
 datamodule:
   module_type: "MultitaskFromSmilesDataModule"
   args: # Matches that in the test_multitask_datamodule.py case.
@@ -194,7 +166,7 @@ predictor:
     l1000_vcap: []
   loss_fun:
     l1000_vcap:
-      name: hybrid_ce_ipu
+      name: hybrid_ce
       n_brackets: 5
   random_seed: *seed
   optim_kwargs:
@@ -216,13 +188,13 @@ predictor:
 metrics:
   l1000_vcap: &classif_metrics
     - name: auroc
-      metric: auroc_ipu
+      metric: auroc
       num_classes: 5
       task: multiclass
       multitask_handling: mean-per-label
       threshold_kwargs: null
     - name: avpr
-      metric: average_precision_ipu
+      metric: average_precision
       num_classes: 5
       task: multiclass
       target_to_int: True
diff --git a/expts/neurips2023_configs/single_task_gine/config_large_gine_g25.yaml b/expts/neurips2023_configs/single_task_gine/config_large_gine_g25.yaml
index 02679153c..345620aff 100644
--- a/expts/neurips2023_configs/single_task_gine/config_large_gine_g25.yaml
+++ b/expts/neurips2023_configs/single_task_gine/config_large_gine_g25.yaml
@@ -1,22 +1,14 @@
-# Testing the gine model with the PCQMv2 dataset on IPU.
+# Testing the gine model with the PCQMv2 dataset.
 constants:
   name: &name neurips2023_large_data_gine_g25
   seed: &seed 42
   raise_train_error: true   # Whether the code should raise an error if it crashes during training
 
 accelerator:
-  type: ipu  # cpu or ipu or gpu
+  type: gpu  # cpu or gpu
   config_override:
     datamodule:
       args:
-        ipu_dataloader_training_opts:
-          mode: async
-          max_num_nodes_per_graph: 30 # train max nodes: 20, max_edges: 54
-          max_num_edges_per_graph: 100
-        ipu_dataloader_inference_opts:
-          mode: async
-          max_num_nodes_per_graph: 30 # valid max nodes: 51, max_edges: 118
-          max_num_edges_per_graph: 100
         # Data handling-related
         batch_size_training: 10
         batch_size_inference: 10
@@ -28,26 +20,6 @@ accelerator:
         precision: 32
         accumulate_grad_batches: 8
 
-  ipu_config:
-    - deviceIterations(10) # IPU would require large batches to be ready for the model.
-    - replicationFactor(16)
-    # - enableProfiling("graph_analyser")       # The folder where the profile will be stored
-    # - enableExecutableCaching("pop_compiler_cache")
-    - TensorLocations.numIOTiles(128)
-    - _Popart.set("defaultBufferingDepth", 128)
-    - Precision.enableStochasticRounding(True)
-
-# accelerator:
-#   type: cpu  # cpu or ipu or gpu
-#   config_override:
-#     datamodule:
-#       batch_size_training: 64
-#       batch_size_inference: 256
-#     trainer:
-#       trainer:
-#         precision: 32
-#         accumulate_grad_batches: 1
-
 datamodule:
   module_type: "MultitaskFromSmilesDataModule"
   args: # Matches that in the test_multitask_datamodule.py case.
@@ -204,7 +176,7 @@ predictor:
   metrics_on_training_set:
     pcqm4m_g25: []
   loss_fun:
-    pcqm4m_g25: mae_ipu
+    pcqm4m_g25: mae
   random_seed: *seed
   optim_kwargs:
     lr: 1.e-4 # warmup can be scheduled using torch_scheduler_kwargs
@@ -225,17 +197,17 @@ predictor:
 metrics:
   pcqm4m_g25: &pcqm_metrics
     - name: mae
-      metric: mae_ipu
+      metric: mae
       target_nan_mask: null
       multitask_handling: flatten
       threshold_kwargs: null
     - name: pearsonr
-      metric: pearsonr_ipu
+      metric: pearsonr
       threshold_kwargs: null
       target_nan_mask: null
       multitask_handling: mean-per-label
     - name: r2
-      metric: r2_score_ipu
+      metric: r2_score
       threshold_kwargs: null
       target_nan_mask: null
       multitask_handling: mean-per-label
diff --git a/expts/neurips2023_configs/single_task_gine/config_large_gine_mcf7.yaml b/expts/neurips2023_configs/single_task_gine/config_large_gine_mcf7.yaml
index 0506dbfea..c7a03a80f 100644
--- a/expts/neurips2023_configs/single_task_gine/config_large_gine_mcf7.yaml
+++ b/expts/neurips2023_configs/single_task_gine/config_large_gine_mcf7.yaml
@@ -1,22 +1,14 @@
-# Testing the gine model with the PCQMv2 dataset on IPU.
+# Testing the gine model with the PCQMv2 dataset.
 constants:
   name: &name neurips2023_large_data_gine_mcf7
   seed: &seed 42
   raise_train_error: true   # Whether the code should raise an error if it crashes during training
 
 accelerator:
-  type: ipu  # cpu or ipu or gpu
+  type: gpu  # cpu or gpu
   config_override:
     datamodule:
       args:
-        ipu_dataloader_training_opts:
-          mode: async
-          max_num_nodes_per_graph: 60 # train max nodes: 20, max_edges: 54
-          max_num_edges_per_graph: 100
-        ipu_dataloader_inference_opts:
-          mode: async
-          max_num_nodes_per_graph: 60 # valid max nodes: 51, max_edges: 118
-          max_num_edges_per_graph: 100
         # Data handling-related
         batch_size_training: 10
         batch_size_inference: 2
@@ -28,26 +20,6 @@ accelerator:
         precision: 32
         accumulate_grad_batches: 8
 
-  ipu_config:
-    - deviceIterations(1) # IPU would require large batches to be ready for the model.
-    - replicationFactor(16)
-    # - enableProfiling("graph_analyser")       # The folder where the profile will be stored
-    # - enableExecutableCaching("pop_compiler_cache")
-    - TensorLocations.numIOTiles(128)
-    - _Popart.set("defaultBufferingDepth", 128)
-    - Precision.enableStochasticRounding(True)
-
-# accelerator:
-#   type: cpu  # cpu or ipu or gpu
-#   config_override:
-#     datamodule:
-#       batch_size_training: 64
-#       batch_size_inference: 256
-#     trainer:
-#       trainer:
-#         precision: 32
-#         accumulate_grad_batches: 1
-
 datamodule:
   module_type: "MultitaskFromSmilesDataModule"
   args: # Matches that in the test_multitask_datamodule.py case.
@@ -204,7 +176,7 @@ predictor:
     l1000_mcf7: []
   loss_fun:
     l1000_mcf7:
-      name: hybrid_ce_ipu
+      name: hybrid_ce
       n_brackets: 5
   random_seed: *seed
   optim_kwargs:
@@ -226,13 +198,13 @@ predictor:
 metrics:
   l1000_mcf7: &classif_metrics
     - name: auroc
-      metric: auroc_ipu
+      metric: auroc
       num_classes: 5
       task: multiclass
       multitask_handling: mean-per-label
       threshold_kwargs: null
     - name: avpr
-      metric: average_precision_ipu
+      metric: average_precision
       num_classes: 5
       task: multiclass
       target_to_int: True
diff --git a/expts/neurips2023_configs/single_task_gine/config_large_gine_n4.yaml b/expts/neurips2023_configs/single_task_gine/config_large_gine_n4.yaml
index 58bad3bbc..edba240f9 100644
--- a/expts/neurips2023_configs/single_task_gine/config_large_gine_n4.yaml
+++ b/expts/neurips2023_configs/single_task_gine/config_large_gine_n4.yaml
@@ -1,22 +1,14 @@
-# Testing the gine model with the PCQMv2 dataset on IPU.
+# Testing the gine model with the PCQMv2 dataset.
 constants:
   name: &name neurips2023_large_data_gine_n4
   seed: &seed 42
   raise_train_error: true   # Whether the code should raise an error if it crashes during training
 
 accelerator:
-  type: ipu  # cpu or ipu or gpu
+  type: gpu  # cpu or gpu
   config_override:
     datamodule:
       args:
-        ipu_dataloader_training_opts:
-          mode: async
-          max_num_nodes_per_graph: 30 # train max nodes: 20, max_edges: 54
-          max_num_edges_per_graph: 100
-        ipu_dataloader_inference_opts:
-          mode: async
-          max_num_nodes_per_graph: 30 # valid max nodes: 51, max_edges: 118
-          max_num_edges_per_graph: 100
         # Data handling-related
         batch_size_training: 10
         batch_size_inference: 10
@@ -28,26 +20,6 @@ accelerator:
         precision: 32
         accumulate_grad_batches: 8
 
-  ipu_config:
-    - deviceIterations(10) # IPU would require large batches to be ready for the model.
-    - replicationFactor(16)
-    # - enableProfiling("graph_analyser")       # The folder where the profile will be stored
-    # - enableExecutableCaching("pop_compiler_cache")
-    - TensorLocations.numIOTiles(128)
-    - _Popart.set("defaultBufferingDepth", 128)
-    - Precision.enableStochasticRounding(True)
-
-# accelerator:
-#   type: cpu  # cpu or ipu or gpu
-#   config_override:
-#     datamodule:
-#       batch_size_training: 64
-#       batch_size_inference: 256
-#     trainer:
-#       trainer:
-#         precision: 32
-#         accumulate_grad_batches: 1
-
 datamodule:
   module_type: "MultitaskFromSmilesDataModule"
   args: # Matches that in the test_multitask_datamodule.py case.
@@ -205,7 +177,7 @@ predictor:
   metrics_on_training_set:
     pcqm4m_n4: []
   loss_fun:
-    pcqm4m_n4: mae_ipu
+    pcqm4m_n4: mae
   random_seed: *seed
   optim_kwargs:
     lr: 1.e-4 # warmup can be scheduled using torch_scheduler_kwargs
@@ -226,17 +198,17 @@ predictor:
 metrics:
   pcqm4m_n4: &pcqm_metrics
     - name: mae
-      metric: mae_ipu
+      metric: mae
       target_nan_mask: null
       multitask_handling: flatten
       threshold_kwargs: null
     - name: pearsonr
-      metric: pearsonr_ipu
+      metric: pearsonr
       threshold_kwargs: null
       target_nan_mask: null
       multitask_handling: mean-per-label
     - name: r2
-      metric: r2_score_ipu
+      metric: r2_score
       threshold_kwargs: null
       target_nan_mask: null
       multitask_handling: mean-per-label
diff --git a/expts/neurips2023_configs/single_task_gine/config_large_gine_pcba.yaml b/expts/neurips2023_configs/single_task_gine/config_large_gine_pcba.yaml
index 3ce9ffde2..32c1af644 100644
--- a/expts/neurips2023_configs/single_task_gine/config_large_gine_pcba.yaml
+++ b/expts/neurips2023_configs/single_task_gine/config_large_gine_pcba.yaml
@@ -1,22 +1,14 @@
-# Testing the gine model with the PCQMv2 dataset on IPU.
+# Testing the gine model with the PCQMv2 dataset.
 constants:
   name: &name neurips2023_large_data_gine_pcba
   seed: &seed 42
   raise_train_error: true   # Whether the code should raise an error if it crashes during training
 
 accelerator:
-  type: ipu  # cpu or ipu or gpu
+  type: gpu  # cpu or gpu
   config_override:
     datamodule:
       args:
-        ipu_dataloader_training_opts:
-          mode: async
-          max_num_nodes_per_graph: 60 # train max nodes: 20, max_edges: 54
-          max_num_edges_per_graph: 100
-        ipu_dataloader_inference_opts:
-          mode: async
-          max_num_nodes_per_graph: 200 # valid max nodes: 51, max_edges: 118
-          max_num_edges_per_graph: 400
         # Data handling-related
         batch_size_training: 10
         batch_size_inference: 2
@@ -28,26 +20,6 @@ accelerator:
         precision: 32
         accumulate_grad_batches: 8
 
-  ipu_config:
-    - deviceIterations(10) # IPU would require large batches to be ready for the model.
-    - replicationFactor(16)
-    # - enableProfiling("graph_analyser")       # The folder where the profile will be stored
-    # - enableExecutableCaching("pop_compiler_cache")
-    - TensorLocations.numIOTiles(128)
-    - _Popart.set("defaultBufferingDepth", 128)
-    - Precision.enableStochasticRounding(True)
-
-# accelerator:
-#   type: cpu  # cpu or ipu or gpu
-#   config_override:
-#     datamodule:
-#       batch_size_training: 64
-#       batch_size_inference: 256
-#     trainer:
-#       trainer:
-#         precision: 32
-#         accumulate_grad_batches: 1
-
 datamodule:
   module_type: "MultitaskFromSmilesDataModule"
   args: # Matches that in the test_multitask_datamodule.py case.
@@ -201,7 +173,7 @@ predictor:
   metrics_on_training_set:
     pcba_1328: []
   loss_fun:
-    pcba_1328: bce_ipu
+    pcba_1328: bce
   random_seed: *seed
   optim_kwargs:
     lr: 1.e-4 # warmup can be scheduled using torch_scheduler_kwargs
@@ -222,12 +194,12 @@ predictor:
 metrics:
   pcba_1328:
     - name: auroc
-      metric: auroc_ipu
+      metric: auroc
       task: binary
       multitask_handling: mean-per-label
       threshold_kwargs: null
     - name: avpr
-      metric: average_precision_ipu
+      metric: average_precision
       task: binary
       multitask_handling: mean-per-label
       threshold_kwargs: null
diff --git a/expts/neurips2023_configs/single_task_gine/config_large_gine_pcq.yaml b/expts/neurips2023_configs/single_task_gine/config_large_gine_pcq.yaml
index d541b9b04..15026ae74 100644
--- a/expts/neurips2023_configs/single_task_gine/config_large_gine_pcq.yaml
+++ b/expts/neurips2023_configs/single_task_gine/config_large_gine_pcq.yaml
@@ -1,22 +1,14 @@
-# Testing the gine model with the PCQMv2 dataset on IPU.
+# Testing the gine model with the PCQMv2 dataset.
 constants:
   name: &name neurips2023_large_data_gine_pcq
   seed: &seed 42
   raise_train_error: true   # Whether the code should raise an error if it crashes during training
 
 accelerator:
-  type: ipu  # cpu or ipu or gpu
+  type: gpu  # cpu or gpu
   config_override:
     datamodule:
       args:
-        ipu_dataloader_training_opts:
-          mode: async
-          max_num_nodes_per_graph: 30 # train max nodes: 20, max_edges: 54
-          max_num_edges_per_graph: 100
-        ipu_dataloader_inference_opts:
-          mode: async
-          max_num_nodes_per_graph: 30 # valid max nodes: 51, max_edges: 118
-          max_num_edges_per_graph: 100
         # Data handling-related
         batch_size_training: 10
         batch_size_inference: 10
@@ -28,26 +20,6 @@ accelerator:
         precision: 32
         accumulate_grad_batches: 8
 
-  ipu_config:
-    - deviceIterations(10) # IPU would require large batches to be ready for the model.
-    - replicationFactor(16)
-    # - enableProfiling("graph_analyser")       # The folder where the profile will be stored
-    # - enableExecutableCaching("pop_compiler_cache")
-    - TensorLocations.numIOTiles(128)
-    - _Popart.set("defaultBufferingDepth", 128)
-    - Precision.enableStochasticRounding(True)
-
-# accelerator:
-#   type: cpu  # cpu or ipu or gpu
-#   config_override:
-#     datamodule:
-#       batch_size_training: 64
-#       batch_size_inference: 256
-#     trainer:
-#       trainer:
-#         precision: 32
-#         accumulate_grad_batches: 1
-
 datamodule:
   module_type: "MultitaskFromSmilesDataModule"
   args: # Matches that in the test_multitask_datamodule.py case.
@@ -243,8 +215,8 @@ predictor:
     pcqm4m_g25: []
     pcqm4m_n4: []
   loss_fun:
-    pcqm4m_g25: mae_ipu
-    pcqm4m_n4: mae_ipu
+    pcqm4m_g25: mae
+    pcqm4m_n4: mae
   random_seed: *seed
   optim_kwargs:
     lr: 1.e-4 # warmup can be scheduled using torch_scheduler_kwargs
@@ -265,17 +237,17 @@ predictor:
 metrics:
   pcqm4m_g25: &pcqm_metrics
     - name: mae
-      metric: mae_ipu
+      metric: mae
       target_nan_mask: null
       multitask_handling: flatten
       threshold_kwargs: null
     - name: pearsonr
-      metric: pearsonr_ipu
+      metric: pearsonr
       threshold_kwargs: null
       target_nan_mask: null
       multitask_handling: mean-per-label
     - name: r2
-      metric: r2_score_ipu
+      metric: r2_score
       threshold_kwargs: null
       target_nan_mask: null
       multitask_handling: mean-per-label
diff --git a/expts/neurips2023_configs/single_task_gine/config_large_gine_vcap.yaml b/expts/neurips2023_configs/single_task_gine/config_large_gine_vcap.yaml
index 121d74ddb..089aa8ed3 100644
--- a/expts/neurips2023_configs/single_task_gine/config_large_gine_vcap.yaml
+++ b/expts/neurips2023_configs/single_task_gine/config_large_gine_vcap.yaml
@@ -1,22 +1,14 @@
-# Testing the gine model with the PCQMv2 dataset on IPU.
+# Testing the gine model with the PCQMv2 dataset.
 constants:
   name: &name neurips2023_large_data_gine_vcap
   seed: &seed 42
   raise_train_error: true   # Whether the code should raise an error if it crashes during training
 
 accelerator:
-  type: ipu  # cpu or ipu or gpu
+  type: gpu  # cpu or gpu
   config_override:
     datamodule:
       args:
-        ipu_dataloader_training_opts:
-          mode: async
-          max_num_nodes_per_graph: 60 # train max nodes: 20, max_edges: 54
-          max_num_edges_per_graph: 100
-        ipu_dataloader_inference_opts:
-          mode: async
-          max_num_nodes_per_graph: 60 # valid max nodes: 51, max_edges: 118
-          max_num_edges_per_graph: 150
         # Data handling-related
         batch_size_training: 10
         batch_size_inference: 2
@@ -28,26 +20,6 @@ accelerator:
         precision: 32
         accumulate_grad_batches: 8
 
-  ipu_config:
-    - deviceIterations(1) # IPU would require large batches to be ready for the model.
-    - replicationFactor(16)
-    # - enableProfiling("graph_analyser")       # The folder where the profile will be stored
-    # - enableExecutableCaching("pop_compiler_cache")
-    - TensorLocations.numIOTiles(128)
-    - _Popart.set("defaultBufferingDepth", 128)
-    - Precision.enableStochasticRounding(True)
-
-# accelerator:
-#   type: cpu  # cpu or ipu or gpu
-#   config_override:
-#     datamodule:
-#       batch_size_training: 64
-#       batch_size_inference: 256
-#     trainer:
-#       trainer:
-#         precision: 32
-#         accumulate_grad_batches: 1
-
 datamodule:
   module_type: "MultitaskFromSmilesDataModule"
   args: # Matches that in the test_multitask_datamodule.py case.
@@ -202,7 +174,7 @@ predictor:
     l1000_vcap: []
   loss_fun:
     l1000_vcap:
-      name: hybrid_ce_ipu
+      name: hybrid_ce
       n_brackets: 5
   random_seed: *seed
   optim_kwargs:
@@ -224,13 +196,13 @@ predictor:
 metrics:
   l1000_vcap: &classif_metrics
     - name: auroc
-      metric: auroc_ipu
+      metric: auroc
       num_classes: 5
       task: multiclass
       multitask_handling: mean-per-label
       threshold_kwargs: null
     - name: avpr
-      metric: average_precision_ipu
+      metric: average_precision
       num_classes: 5
       task: multiclass
       target_to_int: True
diff --git a/graphium/config/_loader.py b/graphium/config/_loader.py
index 9596e2eac..7550c6cbb 100644
--- a/graphium/config/_loader.py
+++ b/graphium/config/_loader.py
@@ -33,8 +33,6 @@
 
 from graphium.data.datamodule import BaseDataModule, MultitaskFromSmilesDataModule
 from graphium.finetuning.finetuning_architecture import FullGraphFinetuningNetwork
-from graphium.ipu.ipu_dataloader import IPUDataloaderOptions
-from graphium.ipu.ipu_utils import import_poptorch, load_ipu_options
 from graphium.nn.architectures import FullGraphMultiTaskNetwork
 from graphium.nn.utils import MupMixin
 from graphium.trainer.metrics import MetricWrapper
@@ -62,51 +60,22 @@ def get_accelerator(
     if (accelerator_type == "gpu") and (not torch.cuda.is_available()):
         raise ValueError(f"GPUs selected, but GPUs are not available on this device")
 
-    # Get the IPU info
-    if accelerator_type == "ipu":
-        poptorch = import_poptorch()
-        if poptorch is None:
-            raise ValueError("IPUs selected, but PopTorch is not available")
-        if not poptorch.ipuHardwareIsAvailable():
-            raise ValueError(
-                "IPUs selected, but no IPU is available/visible on this device. "
-                "If you do have IPUs, please check that the IPUOF_VIPU_API_PARTITION_ID and "
-                "IPUOF_VIPU_API_HOST environment variables are set."
-            )
-
     # Fall on cpu at the end
     if accelerator_type is None:
         accelerator_type = "cpu"
     return accelerator_type
 
 
-def _get_ipu_opts(config: Union[omegaconf.DictConfig, Dict[str, Any]]) -> Tuple[str, str]:
-    r"""
-    Get the paths of the IPU-specific config files from the main YAML config
-    """
-
-    accelerator_options = config["accelerator"]
-    accelerator_type = accelerator_options["type"]
-
-    if accelerator_type != "ipu":
-        return None, None
-    ipu_opts = accelerator_options["ipu_config"]
-    ipu_inference_opts = accelerator_options.get("ipu_inference_config", None)
-
-    return ipu_opts, ipu_inference_opts
-
-
 def load_datamodule(
     config: Union[omegaconf.DictConfig, Dict[str, Any]], accelerator_type: str
 ) -> BaseDataModule:
     """
     Load the datamodule from the specified configurations at the key
     `datamodule: args`.
-    If the accelerator is IPU, load the IPU options as well.
 
     Parameters:
         config: The config file, with key `datamodule: args`
-        accelerator_type: The accelerator type, e.g. "cpu", "gpu", "ipu"
+        accelerator_type: The accelerator type, e.g. "cpu", "gpu"
     Returns:
         datamodule: The datamodule used to process and load the data
     """
@@ -118,53 +87,11 @@ def load_datamodule(
     # Instanciate the datamodule
     module_class = DATAMODULE_DICT[config["datamodule"]["module_type"]]
 
-    if accelerator_type != "ipu":
-        datamodule = module_class(
-            **config["datamodule"]["args"],
-        )
-        return datamodule
-
-    # IPU specific adjustments
-    else:
-        ipu_opts, ipu_inference_opts = _get_ipu_opts(config)
-
-        # Default empty values for the IPU configurations
-        ipu_training_opts = None
-
-        ipu_dataloader_training_opts = cfg_data.pop("ipu_dataloader_training_opts", {})
-        ipu_dataloader_inference_opts = cfg_data.pop("ipu_dataloader_inference_opts", {})
-        ipu_training_opts, ipu_inference_opts = load_ipu_options(
-            ipu_opts=ipu_opts,
-            seed=config["constants"]["seed"],
-            model_name=config["constants"]["name"],
-            gradient_accumulation=config["trainer"]["trainer"].get("accumulate_grad_batches", None),
-            ipu_inference_opts=ipu_inference_opts,
-            precision=config["trainer"]["trainer"].get("precision"),
-        )
-
-        # Define the Dataloader options for the IPU on the training sets
-        bz_train = cfg_data["batch_size_training"]
-        ipu_dataloader_training_opts = IPUDataloaderOptions(
-            batch_size=bz_train, **ipu_dataloader_training_opts
-        )
-        ipu_dataloader_training_opts.set_kwargs()
-
-        # Define the Dataloader options for the IPU on the inference sets
-        bz_test = cfg_data["batch_size_inference"]
-        ipu_dataloader_inference_opts = IPUDataloaderOptions(
-            batch_size=bz_test, **ipu_dataloader_inference_opts
-        )
-        ipu_dataloader_inference_opts.set_kwargs()
-
-        datamodule = module_class(
-            ipu_training_opts=ipu_training_opts,
-            ipu_inference_opts=ipu_inference_opts,
-            ipu_dataloader_training_opts=ipu_dataloader_training_opts,
-            ipu_dataloader_inference_opts=ipu_dataloader_inference_opts,
-            **config["datamodule"]["args"],
-        )
+    datamodule = module_class(
+        **config["datamodule"]["args"],
+    )
+    return datamodule
 
-        return datamodule
 
 
 def load_metrics(config: Union[omegaconf.DictConfig, Dict[str, Any]]) -> Dict[str, MetricWrapper]:
@@ -305,17 +232,12 @@ def load_predictor(
     Defining the predictor module, which handles the training logic from `lightning.LighningModule`
     Parameters:
         model_class: The torch Module containing the main forward function
-        accelerator_type: The accelerator type, e.g. "cpu", "gpu", "ipu"
+        accelerator_type: The accelerator type, e.g. "cpu", "gpu"
     Returns:
         predictor: The predictor module
     """
 
-    if accelerator_type == "ipu":
-        from graphium.ipu.ipu_wrapper import PredictorModuleIPU
-
-        predictor_class = PredictorModuleIPU
-    else:
-        predictor_class = PredictorModule
+    predictor_class = PredictorModule
 
     cfg_pred = dict(deepcopy(config["predictor"]))
     predictor = predictor_class(
@@ -383,42 +305,15 @@ def load_trainer(
     Defining the pytorch-lightning Trainer module.
     Parameters:
         config: The config file, with key `trainer`
-        accelerator_type: The accelerator type, e.g. "cpu", "gpu", "ipu"
+        accelerator_type: The accelerator type, e.g. "cpu", "gpu"
         date_time_suffix: The date and time of the current run. To be used for logging.
     Returns:
         trainer: the trainer module
     """
     cfg_trainer = deepcopy(config["trainer"])
 
-    # Define the IPU plugin if required
-    strategy = cfg_trainer["trainer"].pop("strategy", "auto")
-    if accelerator_type == "ipu":
-        ipu_opts, ipu_inference_opts = _get_ipu_opts(config)
-
-        training_opts, inference_opts = load_ipu_options(
-            ipu_opts=ipu_opts,
-            ipu_inference_opts=ipu_inference_opts,
-            seed=config["constants"]["seed"],
-            model_name=config["constants"]["name"],
-            gradient_accumulation=config["trainer"]["trainer"].get("accumulate_grad_batches", None),
-            precision=config["trainer"]["trainer"].get("precision"),
-        )
-
-        if strategy != "auto":
-            raise ValueError("IPUs selected, but strategy is not set to 'auto'")
-
-        from lightning_graphcore import IPUStrategy
-
-        strategy = IPUStrategy(training_opts=training_opts, inference_opts=inference_opts)
-
     # Get devices
     devices = cfg_trainer["trainer"].pop("devices", 1)
-    if accelerator_type == "ipu":
-        devices = 1  # number of IPUs used is defined in the ipu options files
-
-    # Remove the gradient accumulation from IPUs, since it's handled by the device
-    if accelerator_type == "ipu":
-        cfg_trainer["trainer"].pop("accumulate_grad_batches", None)
 
     # Define the early stopping parameters
     trainer_kwargs = {}
@@ -448,7 +343,6 @@ def load_trainer(
 
     trainer = Trainer(
         detect_anomaly=True,
-        strategy=strategy,
         accelerator=accelerator_type,
         devices=devices,
         callbacks=callbacks,
diff --git a/graphium/config/loc-config_largemix.yaml b/graphium/config/loc-config_largemix.yaml
new file mode 100644
index 000000000..2c96e39ba
--- /dev/null
+++ b/graphium/config/loc-config_largemix.yaml
@@ -0,0 +1,424 @@
+constants:
+  wandb:
+    entity: valencelabs
+    project: graphium3.0
+    name: ${constants.scale}/mpnn/large-no_l1000
+    tags:
+    - mpnn
+    - large
+    - no-l1000
+    - ${constants.scale}
+  data_dir: /home/domix/Gitx/graphium/graphium/data/largemix
+  datacache_path: /home/domix/Gitx/graphium/datacache/largemix
+  scale: 10M
+  max_epochs: 50
+  name: scale_mpnn
+  raise_train_error: true
+  seed: 42
+  variants:
+    1M:
+      mup_scale_factor: 0.27
+      epochs: 50
+      batch_size: 1024
+      accumulate_grad_batches: 1
+      depth: 16
+      # mup_base_path: /rxrx/data/user/frederik.wenkel/outgoing/mup/large-no_l1000/mpnn.yaml
+      train_frac: 1.0
+    3M:
+      mup_scale_factor: 0.505
+      epochs: 50
+      batch_size: 1024
+      accumulate_grad_batches: 1
+      depth: 16
+      # mup_base_path: /rxrx/data/user/frederik.wenkel/outgoing/mup/large-no_l1000/mpnn.yaml
+      train_frac: 1.0
+    10M:
+      mup_scale_factor: null
+      epochs: 50
+      batch_size: 1024
+      accumulate_grad_batches: 1
+      depth: 16
+      mup_base_path: null
+      train_frac: 1.0
+    30M:
+      mup_scale_factor: 1.798
+      epochs: 50
+      batch_size: 1024
+      accumulate_grad_batches: 1
+      depth: 16
+      # mup_base_path: /rxrx/data/user/frederik.wenkel/outgoing/mup/large-no_l1000/mpnn.yaml
+      train_frac: 1.0
+    100M:
+      mup_scale_factor: 3.38
+      epochs: 30
+      batch_size: 512
+      accumulate_grad_batches: 2
+      depth: 16
+      # mup_base_path: /rxrx/data/user/frederik.wenkel/outgoing/mup/large-no_l1000/mpnn.yaml
+      train_frac: 1.0
+    300M:
+      mup_scale_factor: 5.91
+      epochs: 30
+      batch_size: 256
+      accumulate_grad_batches: 1
+      depth: 16
+      # mup_base_path: /rxrx/data/user/frederik.wenkel/outgoing/mup/large-no_l1000/mpnn.yaml
+      train_frac: 1.0
+    1B:
+      mup_scale_factor: 11.0
+      epochs: 20
+      batch_size: 256
+      accumulate_grad_batches: 1
+      depth: 16
+      # mup_base_path: /rxrx/data/user/frederik.wenkel/outgoing/mup/large-no_l1000/mpnn.yaml
+      train_frac: 1.0
+    3B:
+      mup_scale_factor: 18.8
+      epochs: 20
+      batch_size: 128
+      accumulate_grad_batches: 1
+      depth: 16
+      # mup_base_path: /rxrx/data/user/frederik.wenkel/outgoing/mup/large-no_l1000/mpnn.yaml
+      train_frac: 1.0
+    125Mol:
+      mup_scale_factor: 3.38
+      epochs: 30
+      batch_size: 512
+      accumulate_grad_batches: 2
+      depth: 16
+      # mup_base_path: /rxrx/data/user/frederik.wenkel/outgoing/mup/large-no_l1000/mpnn.yaml
+      train_frac: 0.125
+    250Mol:
+      mup_scale_factor: 3.38
+      epochs: 30
+      batch_size: 512
+      accumulate_grad_batches: 2
+      depth: 16
+      # mup_base_path: /rxrx/data/user/frederik.wenkel/outgoing/mup/large-no_l1000/mpnn.yaml
+      train_frac: 0.25
+    500Mol:
+      mup_scale_factor: 3.38
+      epochs: 30
+      batch_size: 512
+      accumulate_grad_batches: 2
+      depth: 16
+      # mup_base_path: /rxrx/data/user/frederik.wenkel/outgoing/mup/large-no_l1000/mpnn.yaml
+      train_frac: 0.5
+  dataset_fraction: 1.0
+accelerator:
+  float32_matmul_precision: medium
+  type: gpu
+architecture:
+  mup_scale_factor: ${constants.variants.${constants.scale}.mup_scale_factor}
+  mup_base_path:  ${constants.variants.${constants.scale}.mup_base_path}
+  pre_nn:
+    activation: relu
+    depth: 2
+    dropout: 0.1
+    hidden_dims: 4
+    last_activation: none
+    last_normalization: layer_norm
+    normalization: layer_norm
+    out_dim: 64
+    residual_type: none
+  pre_nn_edges:
+    activation: relu
+    depth: 2
+    dropout: 0.1
+    hidden_dims: 4
+    last_activation: none
+    last_normalization: layer_norm
+    normalization: layer_norm
+    out_dim: 32
+    residual_type: none
+  gnn:
+    activation: gelu
+    depth: ${constants.variants.${constants.scale}.depth}
+    dropout: 0.1
+    hidden_dims: 4
+    hidden_dims_edges: 2
+    out_dim_edges: 2
+    in_dim: 64
+    last_activation: none
+    last_normalization: layer_norm
+    layer_kwargs:
+      mlp_expansion_ratio: 1
+    layer_type: pyg:mpnnplus
+    normalization: layer_norm
+    out_dim: 4
+    residual_type: simple
+    virtual_node: none
+  graph_output_nn:
+    graph:
+      activation: relu
+      depth: 2
+      dropout: 0.1
+      hidden_dims: 4
+      last_activation: none
+      last_normalization: none
+      normalization: layer_norm
+      out_dim: 653
+      pooling:
+      - sum
+      residual_type: none
+    node:
+      activation: relu
+      depth: 2
+      dropout: 0.1
+      hidden_dims: 4
+      last_activation: none
+      last_normalization: none
+      normalization: layer_norm
+      out_dim: 84
+      pooling:
+      - sum
+      residual_type: none
+  model_type: FullGraphMultiTaskNetwork
+  pe_encoders:
+    encoders:
+      la_pos:
+        dropout: 0.1
+        encoder_type: laplacian_pe
+        first_normalization: none
+        hidden_dim: 2
+        input_keys:
+        - laplacian_eigvec
+        - laplacian_eigval
+        model_type: DeepSet
+        num_layers: 2
+        num_layers_post: 1
+        out_dim: 32
+        output_keys:
+        - feat
+      rw_pos:
+        dropout: 0.1
+        encoder_type: mlp
+        first_normalization: layer_norm
+        hidden_dim: 2
+        input_keys:
+        - rw_return_probs
+        normalization: layer_norm
+        num_layers: 2
+        out_dim: 32
+        output_keys:
+        - feat
+    last_norm: None
+    out_dim: 32
+    pool: sum
+  task_heads:
+    pcba_1328:
+      activation: relu
+      depth: 2
+      dropout: 0.1
+      hidden_dims: 4
+      last_activation: none
+      last_normalization: none
+      normalization: layer_norm
+      out_dim: 448
+      residual_type: none
+      task_level: graph
+    pcqm4m_g25:
+      activation: relu
+      depth: 2
+      dropout: 0.1
+      hidden_dims: 4
+      last_activation: none
+      last_normalization: none
+      normalization: layer_norm
+      out_dim: 25
+      residual_type: none
+      task_level: graph
+    pcqm4m_n4:
+      activation: relu
+      depth: 2
+      dropout: 0.1
+      hidden_dims: 4
+      last_activation: none
+      last_normalization: none
+      normalization: layer_norm
+      out_dim: 4
+      residual_type: none
+      task_level: node
+datamodule:
+  args:
+    batch_size_inference: 1024
+    batch_size_training: ${constants.variants.${constants.scale}.batch_size}
+    train_frac: ${constants.variants.${constants.scale}.train_frac}
+    featurization:
+      add_self_loop: false
+      atom_property_list_float:
+      - degree
+      - formal-charge
+      - radical-electron
+      - aromatic
+      - in-ring
+      atom_property_list_onehot:
+      - atomic-number
+      - group
+      - period
+      - total-valence
+      edge_property_list:
+      - bond-type-onehot
+      - stereo
+      - in-ring
+      explicit_H: false
+      max_num_atoms: 50
+      pos_encoding_as_features:
+        pos_types:
+          lap_eigval:
+            disconnected_comp: true
+            normalization: none
+            num_pos: 8
+            pos_level: node
+            pos_type: laplacian_eigval
+          lap_eigvec:
+            disconnected_comp: true
+            normalization: none
+            num_pos: 8
+            pos_level: node
+            pos_type: laplacian_eigvec
+          rw_pos:
+            ksteps: 16
+            pos_level: node
+            pos_type: rw_return_probs
+      use_bonds_weights: false
+    multiprocessing_context: spawn
+    num_workers: 4
+    persistent_workers: true
+    processed_graph_data_path: ${constants.datacache_path}
+    task_specific_args:
+      pcba_1328:
+        df: null
+        df_path: ${constants.data_dir}/PCBA_1328_1564k.parquet
+        epoch_sampling_fraction: 1
+        sample_size: ${constants.dataset_fraction} # use sample_size for test
+        label_cols: 'assayID-1*'
+        smiles_col: SMILES
+        splits_path: ${constants.data_dir}/pcba_1328_random_splits.pt
+        task_level: graph
+      pcqm4m_g25:
+        df: null
+        df_path: ${constants.data_dir}/PCQM4M_G25_N4.parquet
+        epoch_sampling_fraction: 1
+        sample_size: ${constants.dataset_fraction} # use sample_size for test
+        label_cols: graph_*
+        label_normalization:
+          method: normal
+          normalize_val_test: true
+        smiles_col: ordered_smiles
+        splits_path: ${constants.data_dir}/pcqm4m_g25_n4_random_splits.pt
+        task_level: graph
+      pcqm4m_n4:
+        df: null
+        df_path: ${constants.data_dir}/PCQM4M_G25_N4.parquet
+        epoch_sampling_fraction: 1
+        sample_size: ${constants.dataset_fraction} # use sample_size for test
+        label_cols: node_*
+        label_normalization:
+          method: normal
+          normalize_val_test: true
+        seed: 42
+        smiles_col: ordered_smiles
+        splits_path: ${constants.data_dir}/pcqm4m_g25_n4_random_splits.pt
+        task_level: node
+  module_type: MultitaskFromSmilesDataModule
+metrics:
+  # pcba_1328: []
+  # pcqm4m_g25: []
+  # pcqm4m_n4: []
+  pcba_1328:
+  - metric: auroc
+    multitask_handling: mean-per-label
+    name: auroc
+    target_nan_mask: ignore
+    task: binary
+    threshold_kwargs: null
+    target_to_int: true
+  - metric: averageprecision
+    multitask_handling: mean-per-label
+    name: avpr
+    target_nan_mask: 0
+    task: binary
+    threshold_kwargs: null
+    target_to_int: true
+  pcqm4m_g25:
+  - metric: mae
+    multitask_handling: mean-per-label
+    name: mae
+    target_nan_mask: ignore
+    threshold_kwargs: null
+  - metric: pearsonr
+    multitask_handling: mean-per-label
+    name: pearsonr
+    target_nan_mask: ignore
+    threshold_kwargs: null
+  - metric: r2_score
+    multitask_handling: mean-per-label
+    name: r2
+    target_nan_mask: ignore
+    threshold_kwargs: null
+  pcqm4m_n4:
+  - metric: mae
+    multitask_handling: mean-per-label
+    name: mae
+    target_nan_mask: ignore
+    threshold_kwargs: null
+  - metric: pearsonr
+    multitask_handling: mean-per-label
+    name: pearsonr
+    target_nan_mask: ignore
+    threshold_kwargs: null
+  - metric: r2_score
+    multitask_handling: mean-per-label
+    name: r2
+    target_nan_mask: ignore
+    threshold_kwargs: null
+predictor:
+  loss_fun:
+    pcba_1328: bce_logits
+    pcqm4m_g25: mae
+    pcqm4m_n4: mae
+  metrics_every_n_train_steps: 5
+  metrics_on_progress_bar:
+    pcba_1328: []
+    pcqm4m_g25: []
+    pcqm4m_n4: []
+  metrics_on_training_set:
+    pcba_1328: []
+    pcqm4m_g25: [pearsonr]
+    pcqm4m_n4: [pearsonr]
+  multitask_handling: flatten
+  optim_kwargs:
+    lr: ${eval:"0.003/(((${architecture.gnn.depth}+8)/24)**0.5)"}
+  random_seed: 42
+  scheduler_kwargs: null
+  target_nan_mask: ignore
+  torch_scheduler_kwargs:
+    max_num_epochs: ${constants.max_epochs}
+    module_type: WarmUpLinearLR
+    verbose: false
+    warmup_epochs: 5
+tasks: {}
+trainer:
+  logger:
+    name: ${constants.scale}/mpnn/large-no_l1000
+    project: molgps-pretraining
+    save_dir: logs/molgps-pretraining/large-no_l1000/
+  model_checkpoint:
+    save_last: false
+    save_top_k: -1
+    dirpath: model_checkpoints/graphium3/large-no_l1000/${constants.scale}/mpnn/${constants.seed}/
+    every_n_epochs: 5
+    filename: '{epoch}'
+  seed: ${constants.seed}
+  trainer:
+    check_val_every_n_epoch: 1
+    max_epochs: 50
+    min_epochs: 1
+    precision: 32
+    accumulate_grad_batches: ${constants.variants.${constants.scale}.accumulate_grad_batches}
+    num_sanity_val_steps: 2
+    devices: 1
+    strategy: ddp_find_unused_parameters_true
+    limit_train_batches: 20
+    limit_val_batches: 20
diff --git a/graphium/data/collate.py b/graphium/data/collate.py
index cab3151de..0c1f6ef44 100644
--- a/graphium/data/collate.py
+++ b/graphium/data/collate.py
@@ -22,7 +22,6 @@
 from typing import Union, List, Optional, Dict, Type, Any, Iterable
 from torch_geometric.data import Data, Batch
 
-from graphium.utils.packing import fast_packing, get_pack_sizes, node_to_pack_indices_mask
 from loguru import logger
 from graphium.data.utils import get_keys
 from graphium.data.dataset import torch_enum_to_dtype
@@ -34,7 +33,6 @@ def graphium_collate_fn(
     labels_dtype_dict: Optional[Dict[str, Any]] = None,
     mask_nan: Union[str, float, Type[None]] = "raise",
     do_not_collate_keys: List[str] = [],
-    batch_size_per_pack: Optional[int] = None,
 ) -> Union[Any, Dict[str, Any]]:
     """This collate function is identical to the default
     pytorch collate function but add support for `pyg.data.Data` to batch graphs.
@@ -76,12 +74,6 @@ def graphium_collate_fn(
         do_not_batch_keys:
             Keys to ignore for the collate
 
-        batch_size_per_pack: The number of graphs to pack together.
-            This is useful for using packing with the Transformer.
-            If None, no packing is done.
-            Otherwise, indices are generated to map the nodes to the pack they belong to under the key `"pack_from_node_idx"`,
-            with an additional mask to indicate which nodes are from the same graph under the key `"pack_attn_mask"`.
-
     Returns:
         The batched elements. See `torch.utils.data.dataloader.default_collate`.
     """
@@ -113,7 +105,7 @@ def graphium_collate_fn(
             # If a PyG Graph is provided, use the PyG batching
             elif isinstance(elem[key], Data):
                 pyg_graphs = [d[key] for d in elements]
-                batch[key] = collage_pyg_graph(pyg_graphs, num_nodes, batch_size_per_pack=batch_size_per_pack)
+                batch[key] = collage_pyg_graph(pyg_graphs, num_nodes)
 
             # Ignore the collate for specific keys
             elif key in do_not_collate_keys:
@@ -135,7 +127,7 @@ def graphium_collate_fn(
 
 
 def collage_pyg_graph(
-    pyg_graphs: List[Data], num_nodes: List[int], batch_size_per_pack: Optional[int] = None
+    pyg_graphs: List[Data], num_nodes: List[int],
 ):
     """
     Function to collate pytorch geometric graphs.
@@ -144,8 +136,6 @@ def collage_pyg_graph(
 
     Parameters:
         pyg_graphs: List of PyG graphs
-        batch_size_per_pack: The number of graphs to pack together.
-            This is useful for using packing with the Transformer,
     """
 
     # Calculate maximum number of nodes per graph in current batch
@@ -162,22 +152,6 @@ def collage_pyg_graph(
         # Convert edge index to int64
         pyg_graph.edge_index = pyg_graph.edge_index.to(torch.int64)
 
-    # Apply the packing at the mini-batch level. This is useful for using packing with the Transformer,
-    # especially in the case of the large graphs being much larger than the small graphs.
-    # CAREFUL!!! This changes the order of the graphs in the batch, without changing the order of the labels or other objects.
-    # An error is raised temporarily.
-    if batch_size_per_pack is not None:
-        raise NotImplementedError(
-            "Packing is not yet functional, as it changes the order of the graphs in the batch without changing the label order"
-        )
-        packed_graph_idx = fast_packing(num_nodes, batch_size_per_pack)
-
-        # Get the node to pack indices and the mask
-        pack_from_node_idx, pack_attn_mask = node_to_pack_indices_mask(packed_graph_idx, num_nodes)
-        for pyg_graph in pyg_graphs:
-            pyg_graph.pack_from_node_idx = pack_from_node_idx
-            pyg_graph.pack_attn_mask = pack_attn_mask
-
     return Batch.from_data_list(pyg_graphs)
 
 
diff --git a/graphium/data/datamodule.py b/graphium/data/datamodule.py
index edc5e207d..ff3bcae99 100644
--- a/graphium/data/datamodule.py
+++ b/graphium/data/datamodule.py
@@ -102,7 +102,6 @@ def __init__(
         self,
         batch_size_training: int = 16,
         batch_size_inference: int = 16,
-        batch_size_per_pack: Optional[int] = None,
         num_workers: int = 0,
         pin_memory: bool = True,
         persistent_workers: bool = False,
@@ -125,15 +124,6 @@ def __init__(
 
         self.batch_size_training = batch_size_training
         self.batch_size_inference = batch_size_inference
-        self.batch_size_per_pack = batch_size_per_pack
-        if self.batch_size_per_pack is not None:
-            # Check that batch_size_per_pack is a divisor of batch_size_training and batch_size_inference
-            assert (
-                self.batch_size_training % self.batch_size_per_pack == 0
-            ), f"batch_size_training must be a multiple of batch_size_per_pack, provided batch_size_training={self.batch_size_training}, batch_size_per_pack={self.batch_size_per_pack}"
-            assert (
-                self.batch_size_inference % self.batch_size_per_pack == 0
-            ), f"batch_size_inference must be a multiple of batch_size_per_pack, provided batch_size_inference={self.batch_size_inference}, batch_size_per_pack={self.batch_size_per_pack}"
 
         self.num_workers = num_workers
         self.pin_memory = pin_memory
@@ -203,7 +193,7 @@ def get_collate_fn(self, collate_fn):
         if collate_fn is None:
             # Some values become `inf` when changing data type. `mask_nan` deals with that
             collate_fn = partial(
-                graphium_collate_fn, mask_nan=0, batch_size_per_pack=self.batch_size_per_pack
+                graphium_collate_fn, mask_nan=0,
             )
             collate_fn.__name__ = graphium_collate_fn.__name__
 
@@ -492,12 +482,12 @@ def get_dataloader_kwargs(self, stage: RunningStage, shuffle: bool, **kwargs) ->
         """
         loader_kwargs = {}
 
-        # Get batch size and IPU options for training set
+        # Get batch size for training set
         # if stage in [RunningStage.TRAINING, RunningStage.TUNING]:
         if stage in [RunningStage.TRAINING]:
             loader_kwargs["batch_size"] = self.batch_size_training
 
-        # Get batch size and IPU options for validation / testing sets
+        # Get batch size for validation / testing sets
         elif stage in [RunningStage.VALIDATING, RunningStage.TESTING, RunningStage.PREDICTING]:
             loader_kwargs["batch_size"] = self.batch_size_inference
         else:
@@ -723,64 +713,7 @@ def __init__(
         self.epoch_sampling_fraction = epoch_sampling_fraction
 
 
-class IPUDataModuleModifier:
-    def __init__(
-        self,
-        ipu_inference_opts: Optional["poptorch.Options"] = None,
-        ipu_training_opts: Optional["poptorch.Options"] = None,
-        ipu_dataloader_training_opts: Optional["IPUDataloaderOptions"] = None,
-        ipu_dataloader_inference_opts: Optional["IPUDataloaderOptions"] = None,
-        *args,
-        **kwargs,
-    ) -> None:
-        r"""
-        wrapper functions from the a `DataModule` to support IPU and IPU options To be used in dual inheritance, for example:
-        ```
-        IPUDataModule(BaseDataModule, IPUDataModuleModifier):
-            def __init__(self, **kwargs):
-                BaseDataModule.__init__(self, **kwargs)
-                IPUDataModuleModifier.__init__(self, **kwargs)
-        ```
-
-        Parameters:
-            ipu_inference_opts: Options for the IPU in inference mode. Ignore if not using IPUs
-            ipu_training_opts: Options for the IPU in training mode. Ignore if not using IPUs
-            ipu_dataloader_kwargs_train_val: Options for the dataloader for the IPU. Ignore if not using IPUs
-            ipu_dataloader_kwargs_test: Options for the dataloader for the IPU. Ignore if not using IPUs
-            args: Arguments for the `DataModule`
-            kwargs: Keyword arguments for the `DataModule`
-        """
-        self.ipu_inference_opts = ipu_inference_opts
-        self.ipu_training_opts = ipu_training_opts
-        self.ipu_dataloader_training_opts = ipu_dataloader_training_opts
-        self.ipu_dataloader_inference_opts = ipu_dataloader_inference_opts
-
-    def _dataloader(self, dataset: Dataset, **kwargs) -> "poptorch.DataLoader":
-        """
-        Get a poptorch dataloader for a given dataset
-        Parameters:
-            dataset: The dataset to use
-            kwargs: Keyword arguments for the dataloader
-        Returns:
-            The poptorch dataloader
-        """
-
-        # Use regular Dataloader if no IPUs
-        if ("ipu_options" not in kwargs.keys()) or (kwargs["ipu_options"] is None):
-            raise ValueError(f"No IPU options provided.")
-
-        # Initialize the IPU dataloader
-        from graphium.ipu.ipu_dataloader import create_ipu_dataloader
-
-        loader = create_ipu_dataloader(
-            dataset=dataset,
-            **kwargs,
-        )
-
-        return loader
-
-
-class MultitaskFromSmilesDataModule(BaseDataModule, IPUDataModuleModifier):
+class MultitaskFromSmilesDataModule(BaseDataModule):
     def __init__(
         self,
         task_specific_args: Union[Dict[str, DatasetProcessingParams], Dict[str, Any]],
@@ -788,7 +721,6 @@ def __init__(
         featurization: Optional[Union[Dict[str, Any], omegaconf.DictConfig]] = None,
         batch_size_training: int = 16,
         batch_size_inference: int = 16,
-        batch_size_per_pack: Optional[int] = None,
         num_workers: int = 0,
         pin_memory: bool = True,
         persistent_workers: bool = False,
@@ -836,14 +768,12 @@ def __init__(
             self,
             batch_size_training=batch_size_training,
             batch_size_inference=batch_size_inference,
-            batch_size_per_pack=batch_size_per_pack,
             num_workers=num_workers,
             pin_memory=pin_memory,
             persistent_workers=persistent_workers,
             multiprocessing_context=multiprocessing_context,
             collate_fn=collate_fn,
         )
-        IPUDataModuleModifier.__init__(self, **kwargs)
 
         self._len = None
         self.task_specific_args = task_specific_args
@@ -1304,41 +1234,6 @@ def get_folder_size(self, path):
         # check if the data items are actually saved into the folders
         return sum(os.path.getsize(osp.join(path, f)) for f in os.listdir(path))
 
-    def get_dataloader_kwargs(self, stage: RunningStage, shuffle: bool, **kwargs) -> Dict[str, Any]:
-        """
-        Get the options for the dataloader depending on the current stage.
-
-        Parameters:
-            stage: Whether in Training, Validating, Testing, Sanity-checking, Predicting, or Tuning phase.
-            shuffle: set to ``True`` to have the data reshuffled at every epoch.
-
-        Returns:
-            Arguments to pass to the `DataLoader` during initialization
-        """
-        loader_kwargs = super().get_dataloader_kwargs(stage=stage, shuffle=shuffle, **kwargs)
-
-        # Get batch size and IPU options for training set
-        # if stage in [RunningStage.TRAINING, RunningStage.TUNING]:
-        if stage in [RunningStage.TRAINING]:
-            loader_kwargs["ipu_dataloader_options"] = self.ipu_dataloader_training_opts
-            loader_kwargs["ipu_options"] = self.ipu_training_opts
-
-        # Get batch size and IPU options for validation / testing sets
-        elif stage in [RunningStage.VALIDATING, RunningStage.TESTING, RunningStage.PREDICTING]:
-            loader_kwargs["ipu_dataloader_options"] = self.ipu_dataloader_inference_opts
-            loader_kwargs["ipu_options"] = self.ipu_inference_opts
-        else:
-            raise ValueError(f"Wrong value for `stage`. Provided `{stage}`")
-
-        # Remove the IPU options if not available
-        if loader_kwargs["ipu_options"] is None:
-            loader_kwargs.pop("ipu_options")
-            if loader_kwargs["ipu_dataloader_options"] is not None:
-                logger.warning(
-                    "`ipu_dataloader_options` will be ignored since it is provided without `ipu_options`."
-                )
-            loader_kwargs.pop("ipu_dataloader_options")
-        return loader_kwargs
 
     def get_dataloader(
         self, dataset: Dataset, shuffle: bool, stage: RunningStage
@@ -1365,11 +1260,8 @@ def get_dataloader(
             )
             # turn shuffle off when sampler is used as sampler option is mutually exclusive with shuffle
             kwargs["shuffle"] = False
-        is_ipu = ("ipu_options" in kwargs.keys()) and (kwargs.get("ipu_options") is not None)
-        if is_ipu:
-            loader = IPUDataModuleModifier._dataloader(self, dataset=dataset, sampler=sampler, **kwargs)
-        else:
-            loader = BaseDataModule._dataloader(self, dataset=dataset, sampler=sampler, **kwargs)
+
+        loader = BaseDataModule._dataloader(self, dataset=dataset, sampler=sampler, **kwargs)
 
         return loader
 
@@ -1380,7 +1272,6 @@ def get_collate_fn(self, collate_fn):
                 graphium_collate_fn,
                 mask_nan=0,
                 do_not_collate_keys=["smiles", "mol_ids"],
-                batch_size_per_pack=self.batch_size_per_pack,
             )
             collate_fn.__name__ = graphium_collate_fn.__name__
         return collate_fn
@@ -1769,7 +1660,6 @@ def to_dict(self) -> Dict[str, Any]:
         obj_repr["test_size"] = len(self.test_indices) if self.test_indices is not None else None
         obj_repr["batch_size_training"] = self.batch_size_training
         obj_repr["batch_size_inference"] = self.batch_size_inference
-        obj_repr["batch_size_per_pack"] = self.batch_size_per_pack
         obj_repr["num_node_feats"] = self.num_node_feats
         obj_repr["num_node_feats_with_positional_encoding"] = self.num_node_feats_with_positional_encoding
         obj_repr["num_edge_feats"] = self.num_edge_feats
@@ -1797,7 +1687,6 @@ def __init__(
         featurization: Optional[Union[Dict[str, Any], omegaconf.DictConfig]] = None,
         batch_size_training: int = 16,
         batch_size_inference: int = 16,
-        batch_size_per_pack: Optional[int] = None,
         num_workers: int = 0,
         pin_memory: bool = True,
         persistent_workers: bool = False,
@@ -1868,7 +1757,6 @@ def __init__(
         dm_args["featurization"] = featurization
         dm_args["batch_size_training"] = batch_size_training
         dm_args["batch_size_inference"] = batch_size_inference
-        dm_args["batch_size_per_pack"] = batch_size_per_pack
         dm_args["num_workers"] = num_workers
         dm_args["pin_memory"] = pin_memory
         dm_args["persistent_workers"] = persistent_workers
@@ -2049,7 +1937,6 @@ def __init__(
         featurization: Optional[Union[Dict[str, Any], omegaconf.DictConfig]] = None,
         batch_size_training: int = 16,
         batch_size_inference: int = 16,
-        batch_size_per_pack: Optional[int] = None,
         num_workers: int = 0,
         pin_memory: bool = True,
         persistent_workers: bool = False,
@@ -2110,7 +1997,6 @@ def __init__(
             processed_graph_data_path=processed_graph_data_path,
             batch_size_training=batch_size_training,
             batch_size_inference=batch_size_inference,
-            batch_size_per_pack=batch_size_per_pack,
             num_workers=num_workers,
             pin_memory=pin_memory,
             persistent_workers=persistent_workers,
diff --git a/graphium/data/normalization.py b/graphium/data/normalization.py
index e57a9bcc8..d2e8444b0 100644
--- a/graphium/data/normalization.py
+++ b/graphium/data/normalization.py
@@ -112,13 +112,11 @@ def denormalize(self, input):
             return input
         elif self.method == "normal":
             mean, std = torch.tensor(self.data_mean), torch.tensor(self.data_std)
-            if input.device.type != "ipu":  # Cast to device if not on IPU
-                mean, std = mean.to(input.device), std.to(input.device)
+            mean, std = mean.to(input.device), std.to(input.device)
             return (input * std) + mean
         elif self.method == "unit":
             dmax, dmin = torch.tensor(self.data_max), torch.tensor(self.data_min)
-            if input.device.type != "ipu":  # Cast to device if not on IPU
-                dmax, dmin = dmax.to(input.device), dmin.to(input.device)
+            dmax, dmin = dmax.to(input.device), dmin.to(input.device)
             return input * (dmax - dmin) + dmin
         else:
             raise ValueError(f"normalization method {self.method} not recognised.")
diff --git a/graphium/ipu/README.md b/graphium/ipu/README.md
deleted file mode 100644
index 3f592d7a8..000000000
--- a/graphium/ipu/README.md
+++ /dev/null
@@ -1,15 +0,0 @@
-<div align="center">
-    <img src="../../docs/images/logo-title.png" height="80px">
-    <h3>The Graph Of LIfe Library.</h3>
-</div>
-
-
-## What is in this folder? 
-
-code for IPU acceleration support
-
-- `ipu_dataloader.py`: code for handling dataloader on IPU
-- `ipu_losses.py`: code for computing losses on IPU
-- `ipu_simple_lightning.py`: code for pytorch lightning support on IPU
-- `ipu_utils.py`: utils functions for IPU
-- `ipu_wrapper.py`: wrapper code for IPU support 
\ No newline at end of file
diff --git a/graphium/ipu/__init__.py b/graphium/ipu/__init__.py
deleted file mode 100644
index e69de29bb..000000000
diff --git a/graphium/ipu/ipu_dataloader.py b/graphium/ipu/ipu_dataloader.py
deleted file mode 100644
index 5aa7828f4..000000000
--- a/graphium/ipu/ipu_dataloader.py
+++ /dev/null
@@ -1,434 +0,0 @@
-"""
---------------------------------------------------------------------------------
-Copyright (c) 2023 Valence Labs, Recursion Pharmaceuticals and Graphcore Limited.
-
-Use of this software is subject to the terms and conditions outlined in the LICENSE file.
-Unauthorized modification, distribution, or use is prohibited. Provided 'as is' without
-warranties of any kind.
-
-Valence Labs, Recursion Pharmaceuticals and Graphcore Limited are not liable for any damages arising from its use.
-Refer to the LICENSE file for the full terms and conditions.
---------------------------------------------------------------------------------
-"""
-
-
-from typing import Callable, Iterable, Optional, List, Tuple, Dict, Any, Union
-from copy import deepcopy
-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
-from torch_geometric.transforms import BaseTransform
-
-from graphium.data.utils import get_keys
-from graphium.ipu.ipu_utils import import_poptorch
-from graphium.utils.packing import (
-    fast_packing,
-    hybrid_packing,
-    get_pack_sizes,
-    node_to_pack_indices_mask,
-    estimate_max_pack_node_size,
-)
-
-
-@dataclass
-class IPUDataloaderOptions:
-    r"""
-    This data class stores the arguments necessary to instantiate a model for the Predictor.
-
-    Parameters:
-        model_class:
-            pytorch module used to create a model
-
-        model_kwargs:
-            Key-word arguments used to initialize the model from `model_class`.
-    """
-
-    batch_size: int
-    max_num_nodes: Optional[int] = None
-    max_num_nodes_per_graph: Optional[int] = None
-    max_num_edges: Optional[int] = None
-    max_num_edges_per_graph: Optional[int] = None
-    mode: "poptorch.DataLoaderMode" = "Sync"
-
-    def set_kwargs(self):
-        # Get the maximum number of nodes
-        if self.max_num_nodes is not None:
-            assert (
-                self.max_num_nodes_per_graph is None
-            ), "Cannot use `max_num_nodes` and `max_num_nodes_per_graph` simultaneously"
-        elif self.max_num_nodes_per_graph is not None:
-            assert (
-                self.max_num_nodes is None
-            ), "Cannot use `max_num_nodes` and `max_num_nodes_per_graph` simultaneously"
-            self.max_num_nodes = self.max_num_nodes_per_graph * self.batch_size
-        else:
-            raise ValueError("Must provide either `max_num_nodes` or `max_num_nodes_per_graph`")
-
-        # Get the maximum number of edges
-        if self.max_num_edges is not None:
-            assert (
-                self.max_num_edges_per_graph is None
-            ), "Cannot use `max_num_edges` and `max_num_edges_per_graph` simultaneously"
-        elif self.max_num_edges_per_graph is not None:
-            assert (
-                self.max_num_edges is None
-            ), "Cannot use `max_num_edges` and `max_num_edges_per_graph` simultaneously"
-            self.max_num_edges = self.max_num_edges_per_graph * self.batch_size
-        else:
-            raise ValueError("Must provide either `max_num_nodes` or `max_num_nodes_per_graph`")
-
-        # poptorch mode
-        poptorch = import_poptorch()
-        if isinstance(self.mode, str):
-            if self.mode.lower() == "sync":
-                self.mode = poptorch.DataLoaderMode.Sync
-            elif self.mode.lower() == "async":
-                self.mode = poptorch.DataLoaderMode.Async
-            elif self.mode.lower() == "asyncrebatched":
-                self.mode = poptorch.DataLoaderMode.AsyncRebatched
-            else:
-                raise ValueError(f"`{self.mode}` not a valid parameter.")
-
-
-class CombinedBatchingCollator:
-    """
-    Collator object that manages the combined batch size defined as:
-
-        combined_batch_size = batch_size * device_iterations
-                             * replication_factor * gradient_accumulation
-
-    This is intended to be used in combination with the poptorch.DataLoader
-    """
-
-    def __init__(
-        self,
-        batch_size: int,
-        max_num_nodes: int,
-        max_num_edges: int,
-        dataset_max_nodes_per_graph: int,
-        dataset_max_edges_per_graph: int,
-        collate_fn: Optional[Callable] = None,
-    ):
-        """
-        Parameters:
-            batch_size: mini batch size used by the model
-            max_num_nodes: Maximum number of nodes in the batched padded graph
-            max_num_edges: Maximum number of edges in the batched padded graph
-            dataset_max_nodes_per_graph: Maximum number of nodes per graph in the full dataset
-            dataset_max_edges_per_graph: Maximum number of edges per graph in the full dataset
-            collate_fn: Function used to collate (or batch) the single data or graphs together
-        """
-        super().__init__()
-        self.batch_size = batch_size
-        self.collate_fn = collate_fn
-        self.max_num_nodes = max_num_nodes
-        self.max_num_edges = max_num_edges
-        self.dataset_max_nodes_per_graph = dataset_max_nodes_per_graph
-        self.dataset_max_edges_per_graph = dataset_max_edges_per_graph
-
-    def __call__(
-        self, batch: List[Dict[str, Union[Data, Dict[str, Tensor]]]]
-    ) -> Dict[str, Union[Batch, Dict[str, Tensor], Any]]:
-        """
-        Stack tensors, batch the pyg graphs, and pad each tensor to be same size.
-
-        Parameters:
-            batch: The batch of data, including pyg-graphs `Data` and labels `Dict[str, Tensor]` to be padded
-
-        Returns:
-            out_batch: A dictionary where the graphs are batched and the labels or other Tensors are stacked
-        """
-
-        # Sort the batch such that large graphs are paired with small graphs
-        num_nodes = [b["features"].num_nodes for b in batch]
-        packed_indices = hybrid_packing(num_nodes, batch_size=self.batch_size)
-        packs = [[batch[idx] for idx in pack] for pack in packed_indices]
-
-        # Loop all mini-batches within the global batch
-        all_batches = []
-        for pack in packs:
-            if self.collate_fn != None:
-                local_batch = self.collate_fn(pack)
-
-            transform = Pad(
-                max_num_nodes=self.max_num_nodes,
-                max_num_edges=self.max_num_edges,
-                dataset_max_nodes_per_graph=self.dataset_max_nodes_per_graph,
-                dataset_max_edges_per_graph=self.dataset_max_edges_per_graph,
-            )
-
-            local_batch["features"] = transform(local_batch["features"])
-            local_batch["labels"] = transform(local_batch["labels"])
-            all_batches.append(local_batch)
-
-        out_batch = {}
-
-        # Stack tensors in the first dimension to allow IPUs to differentiate between local and global graph
-        all_keys = get_keys(all_batches[0]["labels"])
-        out_batch["labels"] = {
-            key: torch.stack([this_batch["labels"][key] for this_batch in all_batches], 0) for key in all_keys
-        }
-        out_graphs = [this_batch["features"] for this_batch in all_batches]
-        stacked_features = deepcopy(out_graphs[0])
-        for key, val in out_graphs[0].items():
-            if isinstance(val, torch.Tensor):
-                stacked_features[key] = torch.stack([this_graph[key] for this_graph in out_graphs], dim=0)
-
-        out_batch["features"] = stacked_features
-        for key in all_batches[0].keys():
-            if key not in ("features", "labels"):
-                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
-
-
-def create_ipu_dataloader(
-    dataset: Dataset,
-    ipu_dataloader_options: IPUDataloaderOptions,
-    ipu_options: Optional["poptorch.Options"] = None,
-    batch_size: Optional[int] = 1,
-    collate_fn=None,
-    num_workers: Optional[int] = 0,
-    **kwargs,
-) -> "poptorch.DataLoader":
-    """
-    Creates a poptorch.DataLoader for graph datasets
-    Applies the mini-batching method of concatenating multiple graphs into a
-    single graph with multiple disconnected subgraphs. See:
-    https://pytorch-geometric.readthedocs.io/en/2.0.2/notes/batching.html
-
-    Parameters:
-
-        dataset: The torch_geometric.data.Dataset instance from which to
-            load the graph examples for the IPU.
-        ipu_dataloader_options: The options to initialize the Dataloader for IPU
-        ipu_options: The poptorch.Options used by the
-            poptorch.DataLoader. Will use the default options if not provided.
-        batch_size: How many graph examples to load in each batch
-            (default: 1).
-        collate_fn: The function used to collate batches
-        **kwargs (optional): Additional arguments of :class:`poptorch.DataLoader`.
-
-    Returns:
-        The dataloader
-    """
-    poptorch = import_poptorch()
-
-    if ipu_options is None:
-        # Create IPU default options
-        ipu_options = poptorch.Options()
-
-    # Define the collater function
-    collater = CombinedBatchingCollator(
-        batch_size,
-        collate_fn=collate_fn,
-        max_num_nodes=ipu_dataloader_options.max_num_nodes,
-        max_num_edges=ipu_dataloader_options.max_num_edges,
-        dataset_max_nodes_per_graph=dataset.max_num_nodes_per_graph,
-        dataset_max_edges_per_graph=dataset.max_num_edges_per_graph,
-    )
-
-    # Get the global batch size
-    num_nodes = np.asarray(dataset.num_nodes_list)
-    accum = ipu_options.Training.gradient_accumulation
-    repli = ipu_options._values["replication_factor"]
-    device_iter = ipu_options._values["device_iterations"]
-    combined_batch_size = batch_size * accum * repli * device_iter
-    num_batches = len(dataset) // combined_batch_size
-    num_workers = min(num_batches, num_workers)
-    buffer_size = num_batches // num_workers if num_workers > 0 else None
-    buffer_size = 3 if buffer_size is None else buffer_size
-    async_options = {
-        "sharing_strategy": poptorch.SharingStrategy.ForkServer,
-        "early_preload": True,
-        "buffer_size": buffer_size,
-        "load_indefinitely": True,
-        "miss_sleep_time_in_ms": 0,
-    }
-
-    # Estimate the packing size needed
-    max_pack_size, max_pack_size_per_graph = 0, 0
-    for _ in range(4):
-        this_max_pack_size, this_max_pack_size_per_graph = estimate_max_pack_node_size(
-            num_nodes=num_nodes,
-            batch_size=batch_size,
-            combined_batch_size=combined_batch_size,
-        )
-        max_pack_size = max(max_pack_size, this_max_pack_size)
-        max_pack_size_per_graph = max(max_pack_size_per_graph, this_max_pack_size_per_graph)
-
-    max_num_nodes = collater.max_num_nodes
-    # Log the estimated pack size, with warnings if too big or too small
-    logger.info(
-        f"Estimating pack max_pack_size={max_pack_size} or max_pack_size_per_graph={max_pack_size_per_graph}"
-    )
-    logger.info(f"Provided `max_num_nodes={max_num_nodes}`")
-    if max_pack_size > max_num_nodes - 10:
-        logger.warning(
-            f"The value of `max_num_nodes={max_num_nodes}` seems to be insufficient compared to `max_pack_size={max_pack_size}` and will likely crash"
-        )
-    elif max_pack_size < max_num_nodes - 20:
-        logger.warning(
-            f"The value of `max_num_nodes={max_num_nodes}` seems to be large compared to `max_pack_size={max_pack_size}` and will likely waste memory"
-        )
-
-    return poptorch.DataLoader(
-        options=deepcopy(ipu_options),
-        dataset=dataset,
-        batch_size=batch_size,
-        num_workers=num_workers,
-        collate_fn=collater,
-        async_options=async_options,
-        **kwargs,
-    )
-
-
-class Pad(BaseTransform):
-    """
-    Data transform that applies padding to enforce consistent tensor shapes.
-    """
-
-    def __init__(
-        self,
-        max_num_nodes: int,
-        dataset_max_nodes_per_graph,
-        dataset_max_edges_per_graph,
-        max_num_edges: Optional[int] = None,
-        node_value: float = 0,
-        edge_value: float = 0,
-    ):
-        """
-        Parameters:
-            max_num_nodes: The maximum number of nodes for the total padded graph
-            dataset_max_nodes_per_graph: the maximum number of nodes per graph in the dataset
-            dataset_max_edges_per_graph: the maximum number of edges per graph in the dataset
-            max_num_edges: The maximum number of edges for the total padded graph
-            node_value: Value to add to the node padding
-            edge_value: Value to add to the edge padding
-        """
-        super().__init__()
-        self.max_num_nodes = max_num_nodes
-        self.dataset_max_nodes_per_graph = dataset_max_nodes_per_graph
-        self.dataset_max_edges_per_graph = dataset_max_edges_per_graph
-
-        if max_num_edges:
-            self.max_num_edges = max_num_edges
-        else:
-            # Assume fully connected graph
-            self.max_num_edges = max_num_nodes * (max_num_nodes - 1)
-
-        self.node_value = node_value
-        self.edge_value = edge_value
-
-    def validate(self, data):
-        """
-        Validates that the input graph does not exceed the constraints that:
-
-          * the number of nodes must be <= max_num_nodes
-          * the number of edges must be <= max_num_edges
-
-        Returns:
-            Tuple containing the number nodes and the number of edges
-        """
-        num_nodes = data.num_nodes
-        num_edges = data.num_edges
-
-        assert num_nodes <= self.max_num_nodes, (
-            f"Too many nodes. Graph has {num_nodes} nodes " f"and max_num_nodes is {self.max_num_nodes}."
-        )
-
-        assert num_edges <= self.max_num_edges, (
-            f"Too many edges. Graph has {num_edges} edges defined "
-            f"and max_num_edges is {self.max_num_edges}."
-        )
-
-        return num_nodes, num_edges
-
-    def __call__(self, batch: Batch) -> Batch:
-        return self._call(batch)
-
-    def forward(self, batch: Batch) -> Batch:
-        return self._call(batch)
-
-    def _call(self, batch: Batch) -> Batch:
-        """
-        Pad the batch with a fake graphs that has the desired
-        number of nodes and edges.
-        """
-        num_nodes, num_edges = self.validate(batch)
-        num_pad_nodes = self.max_num_nodes - num_nodes
-        num_pad_edges = self.max_num_edges - num_edges
-        # Create a copy to update with padded features
-        new_batch = deepcopy(batch)
-
-        real_graphs = new_batch.to_data_list()
-
-        for g in real_graphs:
-            g.graph_is_true = torch.tensor([1], dtype=bool)
-            g.node_is_true = torch.full([g.num_nodes], True, dtype=bool)
-            g.edge_is_true = torch.full([g.num_edges], True, dtype=bool)
-
-        # create fake graph with the needed # of nodes and edges
-        fake = Data()
-        fake.num_nodes = num_pad_nodes
-        fake.num_edges = num_pad_edges
-        fake.graph_is_true = torch.tensor([False], dtype=bool)
-        fake.node_is_true = torch.full([num_pad_nodes], False, dtype=bool)
-        fake.edge_is_true = torch.full([num_pad_edges], False, dtype=bool)
-
-        for key, value in real_graphs[0]:
-            if not torch.is_tensor(value):
-                continue
-
-            if key == "graph_is_true" or key == "node_is_true" or key == "edge_is_true":
-                continue
-
-            dim = real_graphs[0].__cat_dim__(key, value)
-            pad_shape = list(value.shape)
-
-            if batch.is_node_attr(key):
-                pad_shape[dim] = num_pad_nodes
-                pad_value = self.node_value
-            elif batch.is_edge_attr(key):
-                pad_shape[dim] = num_pad_edges
-                if key == "edge_index":
-                    # Padding edges are self-loops on the first padding node
-                    pad_value = 0
-                else:
-                    pad_value = self.edge_value
-            # identify graph attributes, pad nan label for the fake graph
-            elif key.startswith("graph_"):
-                num_pad_graphs = 1  # we pad with one big fake graph
-                pad_shape[dim] = num_pad_graphs
-                pad_value = float("nan")
-            else:
-                continue
-
-            pad_value = value.new_full(pad_shape, pad_value)
-            fake[key] = torch.cat([pad_value], dim=dim)
-        real_graphs.append(fake)
-        new_batch = Batch.from_data_list(real_graphs)
-
-        if "num_nodes" in new_batch:
-            new_batch.num_nodes = self.max_num_nodes
-
-        return new_batch
-
-    def __repr__(self) -> str:
-        s = f"{self.__class__.__name__}("
-        s += f"max_num_nodes={self.max_num_nodes}, "
-        s += f"max_num_edges={self.max_num_edges}, "
-        s += f"node_value={self.node_value}, "
-        s += f"edge_value={self.edge_value})"
-        return s
diff --git a/graphium/ipu/ipu_losses.py b/graphium/ipu/ipu_losses.py
deleted file mode 100644
index 6bc434ae4..000000000
--- a/graphium/ipu/ipu_losses.py
+++ /dev/null
@@ -1,196 +0,0 @@
-"""
---------------------------------------------------------------------------------
-Copyright (c) 2023 Valence Labs, Recursion Pharmaceuticals and Graphcore Limited.
-
-Use of this software is subject to the terms and conditions outlined in the LICENSE file.
-Unauthorized modification, distribution, or use is prohibited. Provided 'as is' without
-warranties of any kind.
-
-Valence Labs, Recursion Pharmaceuticals and Graphcore Limited are not liable for any damages arising from its use.
-Refer to the LICENSE file for the full terms and conditions.
---------------------------------------------------------------------------------
-"""
-
-
-import torch
-from torch import Tensor
-from torch.nn import BCELoss, BCEWithLogitsLoss, MSELoss, L1Loss
-from torch._C import _infer_size
-from loguru import logger
-from graphium.trainer.losses import HybridCELoss
-
-
-class BCEWithLogitsLossIPU(BCEWithLogitsLoss):
-    """
-    A modified version of the `torch.nn.BCEWithLogitsLoss` that can ignore NaNs
-    by giving them a weight of `0`. This allows it to work with compilation
-    and IPUs since it doesn't modify the tensor's shape.
-    """
-
-    def forward(self, input: Tensor, target: Tensor) -> Tensor:
-        prev_weight = None
-
-        target = target.clone().to(input.dtype)
-        weight = self.weight
-
-        # Get the original weight matrix. If None, set all weights = 1
-        if weight is not None:
-            prev_weight = self.weight.clone()
-            new_size = _infer_size(target.size(), weight.size())
-            weight = weight.expand(new_size).clone()
-        else:
-            weight = torch.ones(target.shape, dtype=input.dtype, device=input.device)
-
-        # Replace the nan-targets by 0 or 1. Take the value closest to the input.
-        # Give a weight of 0 where there are nan-targets
-        nan_targets = target.isnan()
-        nan_targets_0 = (input < 0.5) & nan_targets
-        nan_targets_1 = (input >= 0.5) & nan_targets
-        target[nan_targets_0] = 0.0
-        target[nan_targets_1] = 1.0
-        weight[nan_targets] = 0.0
-
-        # Compute the loss, and rescale by the number of nan elements
-        self.weight = weight
-        loss = super().forward(input, target)
-
-        num_real_targets = (~nan_targets).sum()
-        factor1 = torch.where(num_real_targets > 0, 1, 0)
-        factor2 = torch.where(num_real_targets > 0, 0, 1)
-        loss = factor1 * loss * nan_targets.numel() / (num_real_targets + factor2)
-
-        # Reset the self.weight to its original value
-        self.weight = prev_weight
-
-        return loss
-
-
-class BCELossIPU(BCELoss):
-    """
-    A modified version of the `torch.nn.BCELoss` that can ignore NaNs
-    by giving them a weight of `0`. This allows it to work with compilation
-    and IPUs since it doesn't modify the tensor's shape.
-    """
-
-    def forward(self, input: Tensor, target: Tensor) -> Tensor:
-        prev_weight = None
-
-        target = target.clone().to(input.dtype)
-        weight = self.weight
-
-        # Get the original weight matrix. If None, set all weights = 1
-        if weight is not None:
-            prev_weight = self.weight.clone()
-            new_size = _infer_size(target.size(), weight.size())
-            weight = weight.expand(new_size).clone()
-        else:
-            weight = torch.ones(target.shape, dtype=input.dtype, device=input.device)
-
-        # Replace the nan-targets by 0 or 1. Take the value closest to the input.
-        # Give a weight of 0 where there are nan-targets
-        nan_targets = target.isnan()
-        nan_targets_0 = (input < 0.5) & nan_targets
-        nan_targets_1 = (input >= 0.5) & nan_targets
-        target[nan_targets_0] = 0.0
-        target[nan_targets_1] = 1.0
-        weight[nan_targets] = 0.0
-
-        # Compute the loss, and rescale by the number of nan elements
-        self.weight = weight
-        loss = super().forward(input, target)
-
-        num_real_targets = (~nan_targets).sum()
-        factor1 = torch.where(num_real_targets > 0, 1, 0)
-        factor2 = torch.where(num_real_targets > 0, 0, 1)
-        loss = factor1 * loss * nan_targets.numel() / (num_real_targets + factor2)
-
-        # Reset the self.weight to its original value
-        self.weight = prev_weight
-
-        return loss
-
-
-class MSELossIPU(MSELoss):
-    """
-    A modified version of the `torch.nn.MSELoss` that can ignore NaNs
-    by giving them the same value for both `input` and `target`.
-    This allows it to work with compilation
-    and IPUs since it doesn't modify the tensor's shape.
-    """
-
-    def forward(self, input: Tensor, target: Tensor) -> Tensor:
-        target = target.clone().to(input.dtype)
-        input = input.clone()
-
-        # Replace the nan-targets in the input/target tensors by 0
-        nan_targets = target.isnan()
-        input[nan_targets] = 0.0
-        target[nan_targets] = 0.0
-
-        # Compute the loss, and rescale by the number of nan elements
-        loss = super().forward(input, target)
-
-        num_real_targets = (~nan_targets).sum()
-        factor1 = torch.where(num_real_targets > 0, 1, 0)
-        factor2 = torch.where(num_real_targets > 0, 0, 1)
-        loss = factor1 * loss * nan_targets.numel() / (num_real_targets + factor2)
-
-        return loss
-
-
-class L1LossIPU(L1Loss):
-    """
-    A modified version of the `torch.nn.L1Loss` that can ignore NaNs
-    by giving them the same value for both `input` and `target`.
-    This allows it to work with compilation
-    and IPUs since it doesn't modify the tensor's shape.
-    """
-
-    def forward(self, input: Tensor, target: Tensor) -> Tensor:
-        target = target.clone().to(input.dtype)
-        input = input.clone()
-
-        # Replace the nan-targets in the input/target tensors by 0
-        nan_targets = target.isnan()
-        input[nan_targets] = 0.0
-        target[nan_targets] = 0.0
-
-        # Compute the loss, and rescale by the number of nan elements
-        loss = super().forward(input, target)
-        num_real_targets = (~nan_targets).sum()
-        factor1 = torch.where(num_real_targets > 0, 1, 0)
-        factor2 = torch.where(num_real_targets > 0, 0, 1)
-        loss = factor1 * loss * nan_targets.numel() / (num_real_targets + factor2)
-
-        return loss
-
-
-class HybridCELossIPU(HybridCELoss):
-    def __init__(
-        self,
-        n_brackets,
-        alpha: float = 0.5,
-    ) -> None:
-        """
-        Parameters:
-            n_brackets: the number of brackets that will be used to group the regression targets.
-                Expected to have the same size as the number of classes in the transformed regression task.
-        """
-        super().__init__(n_brackets=n_brackets, alpha=alpha)
-
-    def forward(self, input: Tensor, target: Tensor) -> Tensor:
-        """
-        Parameters:
-            input: (batch_size x n_classes) tensor of logits predicted for each bracket.
-            target: (batch_size) or (batch_size, 1) tensor of target brackets in {0, 1, ..., self.n_brackets}.
-        """
-
-        target = target.clone().to(input.dtype)
-        input = input.clone()
-
-        # Replace the nan-targets in the input/target tensors by 0
-        nan_targets = target.isnan()
-
-        # Compute the loss, and rescale by the number of nan elements
-        loss = super().forward(input, target, nan_targets)
-        return loss
diff --git a/graphium/ipu/ipu_simple_lightning.py b/graphium/ipu/ipu_simple_lightning.py
deleted file mode 100644
index b2fca086e..000000000
--- a/graphium/ipu/ipu_simple_lightning.py
+++ /dev/null
@@ -1,169 +0,0 @@
-"""
---------------------------------------------------------------------------------
-Copyright (c) 2023 Valence Labs, Recursion Pharmaceuticals and Graphcore Limited.
-
-Use of this software is subject to the terms and conditions outlined in the LICENSE file.
-Unauthorized modification, distribution, or use is prohibited. Provided 'as is' without
-warranties of any kind.
-
-Valence Labs, Recursion Pharmaceuticals and Graphcore Limited are not liable for any damages arising from its use.
-Refer to the LICENSE file for the full terms and conditions.
---------------------------------------------------------------------------------
-"""
-
-
-import lightning
-from lightning_graphcore import IPUStrategy
-from lightning.pytorch.loggers import WandbLogger
-
-import torch
-from torch import nn
-
-import torchvision
-import torchvision.transforms as transforms
-
-import mup
-
-from graphium.nn.base_layers import FCLayer
-from graphium.utils.mup import set_base_shapes
-
-
-ON_IPU = True  # Change this line to run on CPU
-SEED = 42
-
-
-# The simple PyTorch model used in each of these examples
-class SimpleTorchModel(torch.nn.Module):
-    def __init__(self, in_dim, hidden_dim, kernel_size, num_classes):
-        super().__init__()
-        self.in_dim = in_dim
-        self.hidden_dim = hidden_dim
-        self.kernel_size = kernel_size
-        self.num_classes = num_classes
-
-        conv_block = nn.Sequential(
-            nn.Conv2d(in_channels=in_dim, out_channels=hidden_dim, kernel_size=kernel_size),
-            nn.BatchNorm2d(hidden_dim),
-            nn.ReLU(),
-            nn.MaxPool2d(kernel_size),
-            nn.MaxPool2d(kernel_size),
-        )
-
-        self.the_network = nn.Sequential(
-            conv_block,
-            torch.nn.Flatten(),
-            FCLayer(4 * hidden_dim, hidden_dim),
-            FCLayer(hidden_dim, hidden_dim),
-            FCLayer(hidden_dim, num_classes, activation=None, is_readout_layer=True),
-            nn.LogSoftmax(1),
-        )
-
-    def make_mup_base_kwargs(self, divide_factor: float = 2.0):
-        return dict(
-            in_dim=self.in_dim,
-            hidden_dim=round(self.hidden_dim / divide_factor),
-            kernel_size=self.kernel_size,
-            num_classes=self.num_classes,
-        )
-
-    def forward(self, x):
-        return self.the_network(x)
-
-
-# This class shows a minimal lightning example. This example uses our own
-# SimpleTorchModel which is a basic 2 conv, 2 FC torch network. It can be
-# found in simple_torch_model.py.
-class SimpleLightning(lightning.LightningModule):
-    def __init__(self, in_dim, hidden_dim, kernel_size, num_classes, on_ipu):
-        super().__init__()
-        self.model = SimpleTorchModel(
-            in_dim=in_dim, hidden_dim=hidden_dim, kernel_size=kernel_size, num_classes=num_classes
-        )
-        self.on_ipu = on_ipu
-
-    def training_step(self, batch, _):
-        x, label = batch
-        prediction = self.model(x)
-        loss = torch.nn.functional.nll_loss(prediction, label)
-        return loss
-
-    def validation_step(self, batch, _):
-        x, label = batch
-        prediction = self.model(x)
-        preds = torch.argmax(prediction, dim=1)
-        acc = torch.sum(preds == label).float() / len(label)
-        loss = torch.nn.functional.nll_loss(prediction, label)
-        return loss, acc
-
-    # PopTorch doesn't currently support logging within steps. Use the Lightning
-    # callback hooks instead.
-    def on_train_batch_end(self, outputs, batch, batch_idx):
-        self.log("StepLoss", outputs["loss"])
-
-    def validation_epoch_end(self, outputs):
-        loss = [out[0] for out in outputs]
-        self.log("val_loss", torch.stack(loss).mean(), prog_bar=True)
-
-        acc = [out[1] for out in outputs]
-        self.log("val_acc", torch.stack(acc).mean(), prog_bar=True)
-
-    def configure_optimizers(self):
-        adam = torch.optim.Adam
-
-        if self.on_ipu:
-            import poptorch
-
-            adam = poptorch.optim.Adam
-
-        optimizer = mup.MuAdam(self.parameters(), lr=0.01, impl=adam)
-        return optimizer
-
-
-if __name__ == "__main__":
-    torch.manual_seed(SEED)
-
-    # Create the model as usual.
-    predictor = SimpleLightning(in_dim=1, hidden_dim=32, kernel_size=3, num_classes=10, on_ipu=ON_IPU)
-    model = predictor.model
-    base = model.__class__(**model.make_mup_base_kwargs(divide_factor=2))
-    predictor.model = set_base_shapes(model, base, rescale_params=False)
-
-    torch.manual_seed(SEED)
-    # Normal PyTorch dataset.
-    train_set = torchvision.datasets.FashionMNIST(
-        "out/FashionMNIST", train=True, download=True, transform=transforms.Compose([transforms.ToTensor()])
-    )
-    val_set = torchvision.datasets.FashionMNIST(
-        "out/FashionMNIST", train=False, download=True, transform=transforms.Compose([transforms.ToTensor()])
-    )
-
-    # Normal PyTorch dataloader.
-    train_loader = torch.utils.data.DataLoader(train_set, batch_size=16, shuffle=True)
-    val_loader = torch.utils.data.DataLoader(val_set, batch_size=16, shuffle=False)
-
-    torch.manual_seed(SEED)
-
-    ipus = None
-    plugins = None
-    if ON_IPU:
-        import poptorch
-
-        training_opts = poptorch.Options()
-        inference_opts = poptorch.Options()
-
-        # Set the seeds
-        training_opts.randomSeed(SEED)
-        inference_opts.randomSeed(SEED)
-        ipus = 1
-        strategy = IPUStrategy(training_opts=training_opts, inference_opts=inference_opts)
-
-    trainer = lightning.Trainer(
-        logger=WandbLogger(),
-        ipus=ipus,
-        max_epochs=3,
-        log_every_n_steps=1,
-        plugins=plugins,
-    )
-
-    # When fit is called the model will be compiled for IPU and will run on the available IPU devices.
-    trainer.fit(predictor, train_dataloaders=train_loader, val_dataloaders=val_loader)
diff --git a/graphium/ipu/ipu_utils.py b/graphium/ipu/ipu_utils.py
deleted file mode 100644
index c5140ecb5..000000000
--- a/graphium/ipu/ipu_utils.py
+++ /dev/null
@@ -1,162 +0,0 @@
-"""
---------------------------------------------------------------------------------
-Copyright (c) 2023 Valence Labs, Recursion Pharmaceuticals and Graphcore Limited.
-
-Use of this software is subject to the terms and conditions outlined in the LICENSE file.
-Unauthorized modification, distribution, or use is prohibited. Provided 'as is' without
-warranties of any kind.
-
-Valence Labs, Recursion Pharmaceuticals and Graphcore Limited are not liable for any damages arising from its use.
-Refer to the LICENSE file for the full terms and conditions.
---------------------------------------------------------------------------------
-"""
-
-
-import os
-import tempfile
-from datetime import datetime
-from copy import deepcopy
-from types import ModuleType
-from typing import Optional, Tuple, List
-import torch
-
-
-def import_poptorch(raise_error=True) -> Optional[ModuleType]:
-    """
-    Import poptorch and returns it.
-    It is wrapped in a function to avoid breaking the code
-    for non-IPU devices which did not install poptorch.
-
-    Parameters:
-        raise_error: Whether to raise an error if poptorch is unavailable.
-            If `False`, return `None`
-
-    Returns:
-        The poptorch module
-
-    """
-    try:
-        import poptorch
-
-        return poptorch
-    except ImportError as e:
-        if raise_error:
-            raise e
-        return
-
-
-def is_running_on_ipu() -> bool:
-    """
-    Returns whether the current module is running on ipu.
-    Needs to be used in the `forward` or `backward` pass.
-    """
-    poptorch = import_poptorch(raise_error=False)
-    on_ipu = (poptorch is not None) and (poptorch.isRunningOnIpu())
-    return on_ipu
-
-
-def load_ipu_options(
-    ipu_opts: List[str],
-    seed: Optional[int] = None,
-    model_name: Optional[str] = None,
-    gradient_accumulation: Optional[int] = None,
-    precision: Optional[int] = None,
-    ipu_inference_opts: Optional[List[str]] = None,
-) -> Tuple["poptorch.Options", "poptorch.Options"]:
-    """
-    Load the IPU options from the config file.
-
-    Parameters:
-        ipu_cfg: The list  configurations for the IPU, written as a list of strings to make use of `poptorch.Options.loadFromFile`
-
-            write a temporary config gile, and read it. See `Options.loadFromFile`
-            #? see the tutorial for IPU options here
-            # https://github.com/graphcore/tutorials/tree/sdk-release-2.6/tutorials/pytorch/efficient_data_loading
-            #? see the full documentation for ipu options here
-            # https://docs.graphcore.ai/projects/poptorch-user-guide/en/latest/reference.html?highlight=options#poptorch.Options
-
-            ***minibatch size***: The number of samples processed by one simple fwd/bwd pass.
-            = # of samples in a minibatch
-
-            ***device iterations***: A device iteration corresponds to one iteration of the training loop executed on the IPU, starting with data-loading and ending with a weight update.
-            In this simple case, when we set n deviceIterations, the host will prepare n mini-batches in an infeed queue so the IPU can perform efficiently n iterations.
-            = # of minibatches to be processed at a time
-            = # of training / backward pass in this call
-
-            ***gradient accumulation factor***: After each backward pass the gradients are accumulated together for K mini-batches. set K in the argument
-            = # of minibatches to accumulate gradients from
-
-            ***replication factor***: Replication describes the process of running multiple instances of the same model simultaneously on different IPUs to achieve data parallelism.
-            If the model requires N IPUs and the replication factor is M, N x M IPUs will be necessary.
-            = # of times the model is copied to speed up computation, each replica of the model is sent a different subset of the dataset
-
-            ***global batch size***: In a single device iteration, many mini-batches may be processed and the resulting gradients accumulated.
-            We call this total number of samples processed for one optimiser step the global batch size.
-            = total number of samples processed for *one optimiser step*
-            = (minibatch size x Gradient accumulation factor) x Number of replicas
-
-        seed: random seed for the IPU
-        model_name: Name of the model, to be used for ipu profiling
-        ipu_inference_opts: optional IPU configuration overrides for inference.
-            If this is provided, options in this file override those in `ipu_file` for inference.
-
-    Returns:
-
-        training_opts: IPU options for the training set.
-
-        inference_opts: IPU options for inference.
-            It differs from the `training_opts` by enforcing `gradientAccumulation` to 1
-
-    """
-
-    poptorch = import_poptorch()
-    ipu_options = poptorch.Options()
-    ipu_opts_file = ipu_options_list_to_file(ipu_opts)
-    ipu_options.loadFromFile(ipu_opts_file.name)
-    ipu_opts_file.close()
-
-    ipu_options.outputMode(poptorch.OutputMode.All)
-    if seed is not None:
-        ipu_options.randomSeed(seed)
-    if model_name is not None:
-        ipu_options.modelName(f"{model_name}_train")
-    if gradient_accumulation is not None:
-        current = ipu_options.Training.gradient_accumulation
-        assert (current == 1) or (
-            current == gradient_accumulation
-        ), f"Received inconsistent gradient accumulation `{current}` and `{gradient_accumulation}"
-        ipu_options.Training.gradientAccumulation(gradient_accumulation)
-
-    if precision == "16-true":
-        # IPUOptions.loadFromFile currently doesn't support setting half partials, doing it here
-        ipu_options.Precision.setPartialsType(torch.half)
-    training_opts = ipu_options
-
-    # Change the inference options to remove gradient accumulation
-    inference_opts = deepcopy(ipu_options)
-    inference_opts.Training.gradientAccumulation(1)
-    if ipu_inference_opts is not None:
-        ipu_inference_opts_file = ipu_options_list_to_file(ipu_inference_opts)
-        inference_opts.loadFromFile(ipu_inference_opts_file.name)
-        ipu_inference_opts_file.close()
-
-    return training_opts, inference_opts
-
-
-def ipu_options_list_to_file(ipu_opts: Optional[List[str]]) -> tempfile._TemporaryFileWrapper:
-    """
-    Create a temporary file from a list of ipu configs, such that it can be read by `poptorch.Options.loadFromFile`
-
-    Parameters:
-        ipu_opts: The list  configurations for the IPU, written as a list of strings to make use of `poptorch.Options.loadFromFile`
-    Returns:
-        tmp_file: The temporary file of ipu configs
-    """
-    if ipu_opts is None:
-        return
-
-    tmp_file = tempfile.NamedTemporaryFile("w", delete=True)
-    for s in ipu_opts:
-        tmp_file.write(s + "\n")
-    tmp_file.flush()
-    return tmp_file
diff --git a/graphium/ipu/ipu_wrapper.py b/graphium/ipu/ipu_wrapper.py
deleted file mode 100644
index 305326015..000000000
--- a/graphium/ipu/ipu_wrapper.py
+++ /dev/null
@@ -1,235 +0,0 @@
-"""
---------------------------------------------------------------------------------
-Copyright (c) 2023 Valence Labs, Recursion Pharmaceuticals and Graphcore Limited.
-
-Use of this software is subject to the terms and conditions outlined in the LICENSE file.
-Unauthorized modification, distribution, or use is prohibited. Provided 'as is' without
-warranties of any kind.
-
-Valence Labs, Recursion Pharmaceuticals and Graphcore Limited are not liable for any damages arising from its use.
-Refer to the LICENSE file for the full terms and conditions.
---------------------------------------------------------------------------------
-"""
-
-
-from typing import Dict, Any, Optional, Callable, Union, Type, Tuple, Iterable
-
-from torch_geometric.data import Batch
-from torch import Tensor
-from lightning_graphcore import IPUStrategy
-from lightning.pytorch.utilities.types import STEP_OUTPUT
-from lightning.pytorch.trainer.states import RunningStage
-
-from graphium.trainer.predictor import PredictorModule
-from graphium.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
-import functools
-import collections
-from graphium.data.utils import get_keys
-
-poptorch = import_poptorch()
-
-
-class PyGArgsParser(poptorch.ICustomArgParser):
-    """
-    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.
-    """
-
-    @staticmethod
-    def sortedTensorKeys(struct: BaseData) -> Iterable[str]:
-        """
-        Find all the keys that map to a tensor value in struct. The keys
-        are returned in sorted order.
-        """
-        all_keys = sorted(get_keys(struct))
-
-        def isTensor(k: str) -> bool:
-            return isinstance(struct[k], torch.Tensor)
-
-        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]
-
-    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}
-
-        for k in get_keys(original_structure):
-            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):
-    """
-    This class wraps around the `PredictorModule` to make it work with IPU and the `IPUPluginGraphium`.
-    """
-
-    def __init__(self, *args, **kwargs):
-        # Import poptorch in a safe way that will work when working with cpu/gpu
-        self.poptorch = import_poptorch()
-        super().__init__(*args, **kwargs)
-
-    @staticmethod
-    def compute_loss(
-        preds: Dict[str, Tensor],
-        targets: Dict[str, Tensor],
-        weights: Optional[Tensor],
-        loss_fun: Dict[str, Callable],
-        target_nan_mask: Union[Type, str] = "ignore",
-        multitask_handling: Optional[str] = None,
-    ) -> Tuple[Tensor, Dict[str, Tensor]]:
-        return PredictorModule.compute_loss(
-            preds, targets, weights, loss_fun, target_nan_mask, multitask_handling
-        )
-
-    def on_train_batch_end(self, outputs, batch, batch_idx):
-        outputs = self.convert_from_fp16(outputs)
-        outputs["loss"] = outputs["loss"][outputs["loss"] != 0].mean()
-        super().on_train_batch_end(outputs, batch, batch_idx)
-
-    def training_step(self, batch, batch_idx) -> Dict[str, Any]:
-        features, labels = batch["features"], batch["labels"]
-        features, labels = self.squeeze_input_dims(features, labels)
-        dict_input = {"features": features, "labels": labels}
-        step_dict = super().training_step(dict_input)
-
-        loss = step_dict.pop("loss")
-        step_dict["loss"] = self.poptorch.identity_loss(loss, reduction="mean")
-        return step_dict
-
-    def validation_step(self, batch, batch_idx) -> Dict[str, Any]:
-        features, labels = batch["features"], batch["labels"]
-        features, labels = self.squeeze_input_dims(features, labels)
-        dict_input = {"features": features, "labels": labels}
-        step_dict = super().validation_step(dict_input)
-
-        return step_dict
-
-    def test_step(self, batch, batch_idx) -> Dict[str, Any]:
-        # Build a dictionary from the tuples
-        features, labels = batch["features"], batch["labels"]
-        features, labels = self.squeeze_input_dims(features, labels)
-        dict_input = {"features": features, "labels": labels}
-        step_dict = super().test_step(dict_input)
-
-        return step_dict
-
-    def predict_step(self, **inputs) -> Dict[str, Any]:
-        # Build a dictionary from the tuples
-        dict_input = inputs
-        step_dict = super().predict_step(dict_input)
-
-        return step_dict
-
-    def on_validation_batch_end(
-        self, outputs: Any, batch: Any, batch_idx: int, dataloader_idx: int = 0
-    ) -> None:
-        # convert data that will be tracked
-        outputs = self.convert_from_fp16(outputs)
-        super().on_validation_batch_end(outputs, batch, batch_idx, dataloader_idx)
-
-    def evaluation_epoch_end(self, outputs: Any):
-        outputs = self.convert_from_fp16(outputs)
-        super().evaluation_epoch_end(outputs)
-
-    def on_test_batch_end(self, outputs: Any, batch: Any, batch_idx: int, dataloader_idx: int = 0) -> None:
-        outputs = self.convert_from_fp16(outputs)
-        super().on_test_batch_end(outputs, batch, batch_idx, dataloader_idx)
-
-    def configure_optimizers(self, impl=None):
-        if impl is None:
-            dtype = self.precision_to_dtype(self.trainer.precision)
-            impl = functools.partial(
-                self.poptorch.optim.Adam,
-                accum_type=dtype,
-                first_order_momentum_accum_type=dtype,
-                second_order_momentum_accum_type=torch.float,
-            )
-        return super().configure_optimizers(impl=impl)
-
-    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
-
-    def convert_from_fp16(self, data: Any) -> Any:
-        """
-        Converts tensors from FP16 to FP32. Useful to convert the IPU program output data
-        """
-        if isinstance(data, collections.Sequence):
-            for idx in range(len(data)):
-                data[idx] = self.convert_from_fp16(data[idx])
-        elif isinstance(data, collections.Mapping):
-            for key in data:
-                data[key] = self.convert_from_fp16(data[key])
-        elif isinstance(data, torch.Tensor) and data.dtype == torch.float16:
-            data = data.float()
-        return data
-
-    def _convert_features_dtype(self, feats):
-        """
-        Converts features to trainer precision rather than model precision.
-        Necessary to run IPU on FP16.
-        """
-        dtype = self.precision_to_dtype(self.trainer.precision)
-
-        # Convert features to dtype
-        if isinstance(feats, torch.Tensor):
-            feats = feats.to(dtype)
-        elif isinstance(feats, (Data, Batch, dict)):
-            for key, val in feats.items():
-                if isinstance(val, torch.Tensor) and (val.is_floating_point()):
-                    feats[key] = val.to(dtype=dtype)
-        else:
-            raise ValueError(f"Unsupported feats type `{type(feats)}` : {feats}")
-        return feats
-
-    def precision_to_dtype(self, precision):
-        return torch.half if precision == "16-true" else torch.float
-
-    def get_num_graphs(self, data: Batch):
-        """
-        IPU specific method to compute the number of graphs in a Batch,
-        that considers gradient accumulation, multiple IPUs and multiple
-        device iterations. Essential to estimate throughput in graphs/s.
-        """
-        num_graphs = torch.max(data.batch, dim=-1).values
-        num_graphs = torch.sum(num_graphs)
-
-        return num_graphs
diff --git a/graphium/ipu/to_dense_batch.py b/graphium/ipu/to_dense_batch.py
deleted file mode 100644
index 9198ccf3f..000000000
--- a/graphium/ipu/to_dense_batch.py
+++ /dev/null
@@ -1,186 +0,0 @@
-"""
---------------------------------------------------------------------------------
-Copyright (c) 2023 Valence Labs, Recursion Pharmaceuticals and Graphcore Limited.
-
-Use of this software is subject to the terms and conditions outlined in the LICENSE file.
-Unauthorized modification, distribution, or use is prohibited. Provided 'as is' without
-warranties of any kind.
-
-Valence Labs, Recursion Pharmaceuticals and Graphcore Limited are not liable for any damages arising from its use.
-Refer to the LICENSE file for the full terms and conditions.
---------------------------------------------------------------------------------
-"""
-
-
-from typing import Optional, Tuple
-
-import torch
-from torch import Tensor
-from torch_scatter import scatter_add
-
-
-def to_sparse_batch(x: Tensor, mask_idx: Tensor):
-    """
-    Reverse function of `to_dense_batch`
-    """
-    return torch.index_select(x.reshape(-1, x.shape[-1]), 0, mask_idx)
-
-
-def to_sparse_batch_from_packed(x: Tensor, pack_from_node_idx: Tensor):
-    """
-    Reverse function of `to_packed_dense_batch`
-    """
-    return x[pack_from_node_idx[:, 0], pack_from_node_idx[:, 1]]
-
-
-def to_dense_batch(
-    x: Tensor,
-    batch: Optional[Tensor] = None,
-    fill_value: float = 0.0,
-    max_num_nodes_per_graph: Optional[int] = None,
-    batch_size: Optional[int] = None,
-    drop_nodes_last_graph=False,
-) -> Tuple[Tensor, Tensor]:
-    r"""Given a sparse batch of node features
-    :math:`\mathbf{X} \in \mathbb{R}^{(N_1 + \ldots + N_B) \times F}` (with
-    :math:`N_i` indicating the number of nodes in graph :math:`i`), creates a
-    dense node feature tensor
-    :math:`\mathbf{X} \in \mathbb{R}^{B \times N_{\max} \times F}` (with
-    :math:`N_{\max} = \max_i^B N_i`).
-    In addition, a mask of shape :math:`\mathbf{M} \in \{ 0, 1 \}^{B \times
-    N_{\max}}` is returned, holding information about the existence of
-    fake-nodes in the dense representation.
-
-    Parameters:
-        x: Node feature matrix
-            :math:`\mathbf{X} \in \mathbb{R}^{(N_1 + \ldots + N_B) \times F}`.
-        batch: Batch vector
-            :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N`, which assigns each
-            node to a specific example. Must be ordered. (default: :obj:`None`)
-        fill_value: The value for invalid entries in the
-            resulting dense output tensor. (default: :obj:`0`)
-        max_num_nodes_per_graph: The size of the output node dimension.
-            (default: :obj:`None`)
-        batch_size: The batch size. (default: :obj:`None`)
-        drop_nodes_last_graph: Whether to drop the nodes of the last graphs that exceed
-            the `max_num_nodes_per_graph`. Useful when the last graph is a padding.
-
-    :rtype: (:class:`Tensor`, :class:`BoolTensor`)
-    """
-    if batch is None and max_num_nodes_per_graph is None:
-        mask = torch.ones(1, x.size(0), dtype=torch.bool, device=x.device)
-        return x.unsqueeze(0), mask
-
-    if batch is None:
-        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
-    if x.device not in ["ipu", "xla"]:
-        num_nodes = scatter_add(batch.new_ones(x.size(0)), batch, dim=0, dim_size=batch_size)
-    else:
-        # Can't use scatter_add here due to PopTorch bug, will be fixed in SDK 3.3
-        arange = torch.arange(batch_size).unsqueeze(-1)
-        num_nodes = batch.eq(arange).sum(dim=-1)
-    cum_nodes = torch.cat([batch.new_zeros(1), num_nodes.cumsum(dim=0)])
-
-    if max_num_nodes_per_graph is None:  # Must be provided on IPU
-        max_num_nodes_per_graph = int(num_nodes.max())
-
-    idx = torch.arange(batch.size(0), dtype=torch.long, device=x.device)
-    idx = (idx - cum_nodes[batch]) + (batch * max_num_nodes_per_graph)
-
-    size = [batch_size * max_num_nodes_per_graph] + list(x.size())[1:]
-
-    out = x.new_full(size, fill_value)
-
-    ##### CHANGES FROM PYG #####
-
-    # In case the last graph represents padding. Drop the overflowing nodes.
-    if drop_nodes_last_graph:
-        num_nodes = num_nodes[:-1]
-        idx[idx >= size[0]] = size[0] - 1
-
-    # Raise error if num_nodes > max_num_nodes
-    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:])
-
-    # 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.int32, device="cpu")
-        mask = mask.to(x.device)
-        # Can't use mask[idx] here due to PopTorch bug, will be fixed in SDK 3.3
-        # mask[idx] = 1
-        # mask = mask.bool()
-        if drop_nodes_last_graph:
-            num_nodes_with_padding = torch.cat((num_nodes, torch.tensor([0], dtype=torch.int32)), dim=0)
-        else:
-            num_nodes_with_padding = num_nodes
-
-        arange = torch.arange(max_num_nodes_per_graph)
-        mask = num_nodes_with_padding.unsqueeze(-1).gt(arange).flatten()
-
-    else:
-        mask = torch.zeros(mask_sz, dtype=torch.bool, device=x.device)
-        mask[idx] = 1
-
-    ##### END CHANGES FROM PYG #####
-
-    mask = mask.view(batch_size, max_num_nodes_per_graph)
-
-    return out, mask, idx  # Added `idx` as a return
-
-
-def to_packed_dense_batch(
-    x: Tensor,
-    pack_from_node_idx: Tensor,
-    pack_attn_mask: Tensor,
-    fill_value: float = 0.0,
-    max_num_nodes_per_pack: Optional[int] = None,
-) -> Tuple[Tensor, Tensor]:
-    r"""Given a sparse batch of node features
-    :math:`\mathbf{X} \in \mathbb{R}^{(N_1 + \ldots + N_B) \times F}` (with
-    :math:`N_i` indicating the number of nodes in graph :math:`i`), creates a
-    dense node feature tensor
-    :math:`\mathbf{X} \in \mathbb{R}^{B \times N_{\max} \times F}` (with
-    :math:`N_{\max} = \max_i^B N_i`).
-    In addition, a mask of shape :math:`\mathbf{M} \in \{ 0, 1 \}^{B \times
-    N_{\max}}` is returned, holding information about the existence of
-    fake-nodes in the dense representation.
-
-    Parameters: # TODO: Update docstring
-        x: Node feature matrix
-            :math:`\mathbf{X} \in \mathbb{R}^{(N_1 + \ldots + N_B) \times F}`.
-        batch: Batch vector
-            :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N`, which assigns each
-            node to a specific example. Must be ordered. (default: :obj:`None`)
-        fill_value: The value for invalid entries in the
-            resulting dense output tensor. (default: :obj:`0`)
-        max_num_nodes_per_graph: The size of the output node dimension.
-            (default: :obj:`None`)
-        batch_size: The batch size. (default: :obj:`None`)
-        drop_nodes_last_graph: Whether to drop the nodes of the last graphs that exceed
-            the `max_num_nodes_per_graph`. Useful when the last graph is a padding.
-
-    :rtype: (:class:`Tensor`, :class:`BoolTensor`)
-    """
-
-    if max_num_nodes_per_pack is None:  # Must be provided on IPU
-        max_num_nodes_per_pack = pack_attn_mask.shape[-1]
-
-    size = [pack_attn_mask[0], max_num_nodes_per_pack] + list(x.size())[1:]
-
-    out = x.new_full(size, fill_value)
-    out[pack_from_node_idx[:, 0], pack_from_node_idx[:, 1]] = x
-
-    return out
diff --git a/graphium/nn/architectures/global_architectures.py b/graphium/nn/architectures/global_architectures.py
index db28adac4..3ee75984b 100644
--- a/graphium/nn/architectures/global_architectures.py
+++ b/graphium/nn/architectures/global_architectures.py
@@ -14,7 +14,7 @@
 
 from typing import Iterable, List, Dict, Literal, Tuple, Union, Callable, Any, Optional, Type
 from torch_geometric.data import Batch
-from graphium.ipu.to_dense_batch import to_dense_batch
+from torch_geometric.utils import to_dense_batch
 from loguru import logger
 
 # Misc imports
@@ -40,10 +40,6 @@
     ResidualConnectionRandom,
 )
 from graphium.nn.utils import MupMixin
-from graphium.ipu.ipu_utils import import_poptorch, is_running_on_ipu
-
-poptorch = import_poptorch(raise_error=False)
-
 import collections
 
 
@@ -1476,8 +1472,6 @@ def __init__(
 
         if accelerator_kwargs is not None:
             accelerator = accelerator_kwargs["_accelerator"]
-            if accelerator == "ipu":
-                self._apply_ipu_options(accelerator_kwargs)
 
         self._check_bad_arguments()
 
@@ -1529,45 +1523,6 @@ def _check_bad_arguments(self):
                     f"Task heads have graph level tasks {', '.join(graph_level_tasks)}, but pooling is none."
                 )
 
-    def _apply_ipu_options(self, ipu_kwargs):
-        gnn_layers_per_ipu = ipu_kwargs.get("gnn_layers_per_ipu")
-        self._apply_ipu_pipeline_split(gnn_layers_per_ipu)
-
-    def _apply_ipu_pipeline_split(self, gnn_layers_per_ipu):
-        r"""
-        Apply pipeline split from accelerator options if applicable
-        """
-
-        if gnn_layers_per_ipu is None:
-            return
-
-        if not isinstance(gnn_layers_per_ipu, collections.abc.Sequence):
-            raise ValueError("gnn_layers_per_ipu must be a Sequence (e.g. a list)")
-
-        valid_ipu_pipeline_lengths = [1, 2, 4, 8, 16]
-        pipeline_length = len(gnn_layers_per_ipu)
-
-        if pipeline_length not in valid_ipu_pipeline_lengths:
-            raise ValueError(
-                f"Length of gnn_layers_per_ipu must be one of {valid_ipu_pipeline_lengths}, "
-                f"got {gnn_layers_per_ipu} of length {pipeline_length} instead"
-            )
-
-        model_depth = len(self.gnn.layers)
-
-        if sum(gnn_layers_per_ipu) != model_depth:
-            raise ValueError(
-                f"The values in gnn_layers_per_ipu must add up to the depth of the model, "
-                f"got {gnn_layers_per_ipu} with total {sum(gnn_layers_per_ipu)} vs model depth "
-                f"of {model_depth}"
-            )
-
-        begin_block_layer_indices = [sum(gnn_layers_per_ipu[:i]) for i in range(1, pipeline_length)]
-
-        for begin_block_layer_index, ipu_id in zip(begin_block_layer_indices, range(1, pipeline_length)):
-            self.gnn.layers[begin_block_layer_index] = poptorch.BeginBlock(
-                self.gnn.layers[begin_block_layer_index], ipu_id=ipu_id
-            )
 
     def _enable_readout_cache(self, module_filter: Optional[Union[str, List[str]]]):
         """
@@ -1934,7 +1889,6 @@ def forward(self, g: Batch):
                 node_feats=g["feat"],
                 batch=g.batch,
                 max_num_nodes=self.max_num_nodes_per_graph,
-                drop_nodes_last_graph=is_running_on_ipu(),
             )
         # Check if at least one graph-level task is present
         if self.task_level == "graph":
@@ -2024,7 +1978,6 @@ def compute_nodepairs(
         max_num_nodes: int = None,
         fill_value: float = float("nan"),
         batch_size: int = None,
-        drop_nodes_last_graph: bool = False,
     ) -> torch.Tensor:
         r"""
         Vectorized implementation of nodepair-level task:
@@ -2035,19 +1988,16 @@ def compute_nodepairs(
             fill_value: The value for invalid entries in the
                 resulting dense output tensor. (default: :obj:`NaN`)
             batch_size: The batch size. (default: :obj:`None`)
-            drop_nodes_last_graph: Whether to drop the nodes of the last graphs that exceed
-                the `max_num_nodes_per_graph`. Useful when the last graph is a padding.
         Returns:
             result: concatenated node features of shape B * max_num_nodes * 2*h,
             where B is number of graphs, max_num_nodes is the chosen maximum number nodes, and h is the feature dim
         """
-        dense_feat, mask, _ = to_dense_batch(
+        dense_feat, mask = to_dense_batch(
             node_feats,
             batch=batch,
             fill_value=fill_value,
             batch_size=batch_size,
-            max_num_nodes_per_graph=max_num_nodes,
-            drop_nodes_last_graph=drop_nodes_last_graph,
+            max_num_nodes=max_num_nodes,
         )
         n = dense_feat.size(1)
         h_X = dense_feat[:, :, None].repeat(1, 1, n, 1)
diff --git a/graphium/nn/base_graph_layer.py b/graphium/nn/base_graph_layer.py
index 66869f888..0986cb9c6 100644
--- a/graphium/nn/base_graph_layer.py
+++ b/graphium/nn/base_graph_layer.py
@@ -258,14 +258,14 @@ 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)
+        Get the maximum number of nodes per graph. Useful for reshaping a compiled model
         """
         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)
+        Set the maximum number of nodes per graph. Useful for reshaping a compiled model
         """
         if value is not None:
             assert isinstance(value, int) and (
@@ -276,14 +276,14 @@ def max_num_nodes_per_graph(self, value: Optional[int]):
     @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)
+        Get the maximum number of nodes per graph. Useful for reshaping a compiled model
         """
         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)
+        Set the maximum number of nodes per graph. Useful for reshaping a compiled model
         """
         if value is not None:
             assert isinstance(value, int) and (
diff --git a/graphium/nn/base_layers.py b/graphium/nn/base_layers.py
index a88e6ad9b..dafba127f 100644
--- a/graphium/nn/base_layers.py
+++ b/graphium/nn/base_layers.py
@@ -26,7 +26,6 @@
 from mup import set_base_shapes, MuReadout
 from torch.nn.functional import linear
 
-from graphium.ipu.ipu_utils import is_running_on_ipu
 
 SUPPORTED_ACTIVATION_MAP = {
     "ReLU",
@@ -156,7 +155,6 @@ def forward(
         query: Tensor,
         key: Tensor,
         value: Tensor,
-        key_padding_mask: Optional[Tensor] = None,
         attn_bias: Optional[Tensor] = None,
         precision: Optional[str] = "32",
         *args,
@@ -186,15 +184,7 @@ def forward(
         attn_weights = q @ k.transpose(-1, -2)
         # [batch, num_heads, nodes, nodes]
         attn_weights += attn_bias
-        key_padding_mask_value = float("-inf") if precision == "32" else -10000
-        # key_padding_mask: [batch, 1, 1, nodes]
-        if key_padding_mask is not None:
-            masked_attn_weights = attn_weights.masked_fill(
-                key_padding_mask.unsqueeze(1).unsqueeze(2).bool(), # The mask is cast to float somewhere in TransformerEncoder
-                key_padding_mask_value,
-            )
-        else:
-            masked_attn_weights = attn_weights
+        masked_attn_weights = attn_weights
         masked_attn_weights = F.softmax(masked_attn_weights, dim=-1)
         attn_probs = F.dropout(masked_attn_weights, p=self.dropout, training=self.training)
         # [batch, num_heads, nodes, nodes] * [batch, num_heads, nodes, head_size] -> [batch, num_heads, nodes, head_size]
@@ -243,9 +233,7 @@ class MuReadoutGraphium(MuReadout):
     Not quite a drop-in replacement for `mup.MuReadout` - you need to specify
     `base_width`.
 
-    Set `base_width` to width of base model passed to `mup.set_base_shapes`
-    to get same results on IPU and CPU. Should still "work" with any other
-    value, but won't give the same results as CPU
+    Set `base_width` to width of base model passed to `mup.set_base_shapes`.
     """
 
     def __init__(self, in_features, *args, **kwargs):
@@ -725,32 +713,23 @@ def forward(
         Parameters:
             input:  `torch.Tensor[total_num_nodes, hidden]`
             batch: batch attribute of the batch object, batch.batch
-            batch_size: The batch size. Must be provided when working on IPU
+            batch_size: The batch size.
 
         Returns:
             torch.Tensor: `torch.Tensor[total_num_nodes, hidde]`
 
         """
-        on_ipu = is_running_on_ipu()
 
         if self.drop_rate > 0:
             keep_prob = 1 - self.drop_rate
 
             # Parse the batch size
             if batch_size is None:
-                if on_ipu:
-                    raise ValueError(
-                        "When using the IPU the batch size must be "
-                        "provided during compilation instead of determined at runtime"
-                    )
-                else:
-                    batch_size = int(batch_idx.max()) + 1
+                batch_size = int(batch_idx.max()) + 1
 
             # mask shape: [num_graphs, 1]
             mask = input.new_empty(batch_size, 1).bernoulli_(keep_prob)
-            # if on_ipu, the last graph is a padded fake graph
-            if on_ipu:
-                mask[-1] = 0
+
             # using gather to extend mask to [total_num_nodes, 1]
             node_mask = mask[batch_idx]
             if keep_prob == 0:
diff --git a/graphium/nn/encoders/gaussian_kernel_pos_encoder.py b/graphium/nn/encoders/gaussian_kernel_pos_encoder.py
index 44ddd5578..19ff813ac 100644
--- a/graphium/nn/encoders/gaussian_kernel_pos_encoder.py
+++ b/graphium/nn/encoders/gaussian_kernel_pos_encoder.py
@@ -2,7 +2,6 @@
 from torch_geometric.data import Batch
 
 from graphium.nn.pyg_layers.utils import PreprocessPositions
-from graphium.ipu.ipu_utils import is_running_on_ipu
 from graphium.nn.encoders.base_encoder import BaseEncoder
 
 
@@ -116,13 +115,10 @@ def forward(self, batch: Batch, key_prefix: Optional[str] = None) -> Dict[str, A
         """
         input_keys = self.parse_input_keys_with_prefix(key_prefix)
 
-        on_ipu = is_running_on_ipu()
         max_num_nodes_per_graph = None
-        if on_ipu:
-            max_num_nodes_per_graph = self.max_num_nodes_per_graph
 
         attn_bias_3d, node_feature_3d = self.preprocess_3d_positions(
-            batch, max_num_nodes_per_graph, on_ipu, positions_3d_key=input_keys[0]
+            batch, max_num_nodes_per_graph, positions_3d_key=input_keys[0]
         )
 
         # Return `attn_bias_3d` if the key starts with 'nodepair_'
diff --git a/graphium/nn/pyg_layers/gps_pyg.py b/graphium/nn/pyg_layers/gps_pyg.py
index 3d8671c53..24c74a6ec 100644
--- a/graphium/nn/pyg_layers/gps_pyg.py
+++ b/graphium/nn/pyg_layers/gps_pyg.py
@@ -18,6 +18,7 @@
 from torch.nn import Module
 from torch import Tensor
 from torch_geometric.data import Batch
+from torch_geometric.utils import to_dense_batch
 from graphium.nn.base_graph_layer import BaseGraphModule
 from graphium.nn.base_layers import FCLayer, MultiheadAttentionMup, MLP
 from graphium.nn.pyg_layers import (
@@ -29,13 +30,6 @@
 )
 from graphium.data.utils import get_keys
 from graphium.utils.decorators import classproperty
-from graphium.ipu.to_dense_batch import (
-    to_dense_batch,
-    to_sparse_batch,
-    to_packed_dense_batch,
-    to_sparse_batch_from_packed,
-)
-from graphium.ipu.ipu_utils import is_running_on_ipu
 
 PYG_LAYERS_DICT = {
     "pyg:gin": GINConvPyg,
@@ -286,7 +280,7 @@ def forward(self, batch: Batch) -> Batch:
         # MLP block, with skip connection
         feat_mlp = self.mlp(feat)
         # Add the droppath to the output of the MLP
-        batch_size = None if feat.device.type != "ipu" else batch.graph_is_true.shape[0]
+        batch_size = None
         if self.droppath_ffn is not None:
             feat_mlp = self.droppath_ffn(feat_mlp, batch.batch, batch_size)
         feat = feat + feat_mlp
@@ -375,50 +369,27 @@ def _to_dense_batch(
         h: Tensor,
         batch: Batch,
         batch_size: Optional[int] = None,
-        max_num_nodes_per_graph: Optional[int] = None,
-        on_ipu: bool = False,
+        max_num_nodes: Optional[int] = None,
     ) -> Tensor:
         """
         Convert the batch of graphs to a dense batch.
         """
 
-        if self._use_packing(batch):
-            attn_mask = batch.pack_attn_mask
-            key_padding_mask = None
-            idx = batch.pack_from_node_idx
-            h_dense = to_packed_dense_batch(
-                h,
-                pack_from_node_idx=idx,
-                pack_attn_mask=attn_mask,
-                max_num_nodes_per_pack=100,  # TODO: This should be a parameter
-            )
-        else:
-            attn_mask = None
-            h_dense, key_padding_mask, idx = to_dense_batch(
-                h,
-                batch=batch.batch,  # The batch index as a vector that indicates for nodes of which graph it belongs to
-                batch_size=batch_size,
-                max_num_nodes_per_graph=max_num_nodes_per_graph,
-                drop_nodes_last_graph=on_ipu,
-            )
-            key_padding_mask = ~key_padding_mask
-        return h_dense, attn_mask, key_padding_mask, idx
-
-    def _to_sparse_batch(self, batch: Batch, h_dense: Tensor, idx: Tensor) -> Tensor:
+        h_dense, key_padding_mask = to_dense_batch(
+            h,
+            batch=batch.batch,  # The batch index as a vector that indicates for nodes of which graph it belongs to
+            batch_size=batch_size,
+            max_num_nodes=max_num_nodes,
+        )
+        key_padding_mask = ~key_padding_mask
+        return h_dense, key_padding_mask
+
+    def _to_sparse_batch(self, batch: Batch, h_dense: Tensor, mask: torch.BoolTensor) -> Tensor:
         """
         Convert the dense batch back to a sparse batch.
         """
-        if self._use_packing(batch):
-            h = to_sparse_batch_from_packed(
-                h_dense,
-                pack_from_node_idx=idx,
-            )
-        else:
-            h = to_sparse_batch(
-                h_dense,
-                mask_idx=idx,
-            )
-        return h
+        
+        return h_dense[mask]
 
     def _self_attention_block(self, feat: Tensor, feat_in: Tensor, batch: Batch) -> Tensor:
         """
@@ -429,21 +400,17 @@ def _self_attention_block(self, feat: Tensor, feat_in: Tensor, batch: Batch) ->
         """
 
         # Multi-head attention.
-        on_ipu = is_running_on_ipu()
         max_num_nodes_per_graph = None
-        if on_ipu:
-            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 feat.device.type != "ipu" else batch.graph_is_true.shape[0]
+        batch_size = None
 
         # h[num_nodes, hidden_dim] -> h_dense[num_graphs, max_num_nodes, hidden_dim]
-        feat_dense, attn_mask, key_padding_mask, idx = self._to_dense_batch(
+        feat_dense, attn_mask = self._to_dense_batch(
             feat,
             batch=batch,  # The batch index as a vector that indicates for nodes of which graph it belongs to
             batch_size=batch_size,
-            max_num_nodes_per_graph=max_num_nodes_per_graph,
-            on_ipu=on_ipu,
+            max_num_nodes=max_num_nodes_per_graph,
         )
 
         attn_bias = None
@@ -452,11 +419,11 @@ def _self_attention_block(self, feat: Tensor, feat_in: Tensor, batch: Batch) ->
 
         # h_dense[num_graphs, max_num_nodes, hidden_dim] -> feat_attn[num_graphs, max_num_nodes, hidden_dim]
         feat_attn = self._sa_block(
-            feat_dense, attn_bias=attn_bias, attn_mask=attn_mask, key_padding_mask=key_padding_mask
+            feat_dense, attn_bias=attn_bias, attn_mask=attn_mask
         )
 
         # feat_attn[num_graphs, max_num_nodes, hidden_dim] -> feat_attn[num_nodes, hidden_dim]
-        feat_attn = self._to_sparse_batch(batch, feat_attn, idx)
+        feat_attn = self._to_sparse_batch(batch, feat_attn, attn_mask)
 
         # Dropout, residual, norm
         if self.dropout_attn is not None:
diff --git a/graphium/nn/pyg_layers/utils.py b/graphium/nn/pyg_layers/utils.py
index 83dc4f737..169e1e5a7 100644
--- a/graphium/nn/pyg_layers/utils.py
+++ b/graphium/nn/pyg_layers/utils.py
@@ -16,13 +16,13 @@
 import torch
 import torch.nn as nn
 from torch_geometric.data import Batch
+from torch_geometric.utils import to_dense_batch
 from typing import Tuple
 from torch import Tensor
 
 from torch_geometric.typing import SparseTensor
 
 from graphium.nn.base_layers import MLP, get_norm
-from graphium.ipu.to_dense_batch import to_dense_batch, to_sparse_batch
 
 
 class PreprocessPositions(nn.Module):
@@ -74,7 +74,7 @@ def __init__(
         self.node_proj = nn.Linear(self.num_kernel, self.embed_dim)
 
     def forward(
-        self, batch: Batch, max_num_nodes_per_graph: int, on_ipu: bool, positions_3d_key: str
+        self, batch: Batch, max_num_nodes_per_graph: int, positions_3d_key: str
     ) -> Tuple[Tensor, Tensor]:
         r"""
         Inputs:
@@ -82,8 +82,6 @@ def forward(
                 Batch object.
             max_num_nodes_per_graph:
                 Maximum number of nodes per graph.
-            on_ipu:
-                If model rus on IPU.
             positions_3d_key:
                 The key of the pyg graph object that contains the 3D positions.
 
@@ -92,17 +90,14 @@ def forward(
         pos = batch[positions_3d_key]
         if self.first_normalization is not None:
             pos = self.first_normalization(pos)
-        batch_size = None if pos.device.type != "ipu" else batch.graph_is_true.shape[0]
-        # batch_size = None if batch.feat.device.type != "ipu" else batch.graph_is_true.shape[0] #[Andy] batch.feat is only available after passing through layers, not a good attribute to check
+        batch_size = None
         # pos: [batch, nodes, 3]
         # padding_mask: [batch, nodes]
-        # idx: [totoal_nodes]
-        pos, mask, idx = to_dense_batch(
+        pos, mask = to_dense_batch(
             pos,
             batch=batch.batch,
             batch_size=batch_size,
-            max_num_nodes_per_graph=max_num_nodes_per_graph,
-            drop_nodes_last_graph=on_ipu,
+            max_num_nodes=max_num_nodes_per_graph,
         )
         # check nan with the pos from to_dense_batch,
         # and generate mask. 1 for nan, 0 for other values.
@@ -153,7 +148,7 @@ def forward(
         # unsqueezed mask size: [batch, 1, 1] apply on tensor [batch, nodes, embed_dim]
         node_feature.masked_fill_(nan_mask.unsqueeze(1).unsqueeze(2).to(torch.bool), 0.0)
         # [total_nodes, embed_dim]
-        node_feature = to_sparse_batch(node_feature, idx)
+        node_feature = node_feature[mask]
 
         return attn_bias, node_feature
 
diff --git a/graphium/utils/packing.py b/graphium/utils/packing.py
deleted file mode 100644
index 6db6856b1..000000000
--- a/graphium/utils/packing.py
+++ /dev/null
@@ -1,330 +0,0 @@
-"""
---------------------------------------------------------------------------------
-Copyright (c) 2023 Valence Labs, Recursion Pharmaceuticals and Graphcore Limited.
-
-Use of this software is subject to the terms and conditions outlined in the LICENSE file.
-Unauthorized modification, distribution, or use is prohibited. Provided 'as is' without
-warranties of any kind.
-
-Valence Labs, Recursion Pharmaceuticals and Graphcore Limited are not liable for any damages arising from its use.
-Refer to the LICENSE file for the full terms and conditions.
---------------------------------------------------------------------------------
-"""
-
-
-from typing import List, Tuple, Iterable, Optional
-import numpy as np
-import torch
-
-
-class MolPack:
-    """
-    Class that keeps track of the number of atoms and indices that are added
-    to each pack. Useful when doing packing, or other forms of smart batching.
-    A pack is a batch, but with optimized memory consumption.
-    """
-
-    def __init__(self):
-        self.num_nodes = 0
-        self.num_graphs = 0
-        self.average_atom = 0
-        self.indices = []
-
-    def add_mol(self, num_nodes: int, idx: int) -> "MolPack":
-        """
-        Add a molecule and it's index to the batch
-
-        Parameters:
-            num_nodes: Number of atoms of the new molecule
-
-            idx: Index associated to the molecule
-        """
-        self.num_nodes += num_nodes
-        self.num_graphs += 1
-        self.average_atom = self.num_nodes / self.num_graphs
-        self.indices.append(idx)
-        return self
-
-    def expected_atoms(self, remaining_mean_num_nodes: float, batch_size: int) -> float:
-        """
-        Given a desired batch size, and given the remaining mean number of
-        atoms, find the expected number of atoms of the current batch when it is full
-
-        Parameters:
-            remaining_mean_num_nodes: Average number of atoms per molecule
-                left to be sampled and distributed across tasks.
-
-            batch_size: Desired batch size
-
-        Returns:
-            expected_atoms: The expected number of atoms in this batch if we
-                sample randomly the remaining molecules.
-        """
-        return self.num_nodes + ((batch_size - self.num_graphs) * remaining_mean_num_nodes)
-
-    def __repr__(self) -> str:
-        """
-        Print the main attributes of the current class
-        """
-        return f"{self.__class__.__name__}(m: {self.num_graphs},\ta: {self.num_nodes},\tav: {self.average_atom:.1f})"
-
-
-def smart_packing(num_nodes: List[int], batch_size: int) -> List[List[int]]:
-    """
-    Simple and fast algorithm for packing graphs such that each batch has roughly the
-    same number of atoms.
-    Has for-loop scalability issues `O(num_graphs * ipu_batch_size)` = `O(num_graphs^2 / batch_size)`
-
-    Parameters:
-        num_nodes: List of the number of atoms per molecule for the entire global batch.
-            Must be of length `batch_size * ipu_batch_size`.
-
-        batch_size: The batch size per iteration, considering a single device and single
-            forward pass.
-            The global batch size is `batch_size * device_iterations * replication_factor * gradient_accumulation`
-
-    Returns:
-        packed_indices: A list of packs, each containing a list of indices, such that
-            if we collect `num_nodes` from the indices, then each pack has roughly the
-            same total number of atoms.
-    """
-
-    # Sort the list
-    num_nodes = np.asarray(num_nodes)
-    argsort_num_nodes = np.argsort(num_nodes)
-    sorted_num_nodes = num_nodes[argsort_num_nodes]
-    ipu_batch_size = int(len(num_nodes) / batch_size)
-    sorted_num_nodes, initial_num_nodes = (
-        sorted_num_nodes[:-ipu_batch_size],
-        sorted_num_nodes[-ipu_batch_size:],
-    )
-    reverse_cumsum = np.sum(sorted_num_nodes) - np.cumsum(sorted_num_nodes) + sorted_num_nodes[-1]
-
-    # Start with the largest element in separate packs
-    mol_batches = [
-        MolPack().add_mol(initial_num_nodes[-ii - 1], argsort_num_nodes[-ii - 1])
-        for ii in range(ipu_batch_size)
-    ]
-
-    # Loop from smallest to largest molecule, and add each molecule to the pack with smallest expected sum
-    for ii, num_atom in enumerate(sorted_num_nodes):
-        remaining_mean = reverse_cumsum[ii] / (len(sorted_num_nodes) - ii)
-        max_expected, idx_max_expected = 0, 0
-        for jj, m in enumerate(mol_batches):
-            if m.num_graphs >= batch_size:
-                continue
-            expected = m.num_nodes + (
-                (batch_size - m.num_graphs) * remaining_mean
-            )  # Faster than calling m.expected_atoms
-            if expected > max_expected:
-                max_expected = expected
-                idx_max_expected = jj
-        mol_batches[idx_max_expected].add_mol(num_atom, argsort_num_nodes[ii])
-
-    packed_indices = [batch.indices for batch in mol_batches]
-
-    return packed_indices
-
-
-def fast_packing(num_nodes: List[int], batch_size: int) -> List[List[int]]:
-    """
-    Super fast algorithm for packing graphs such that each batch has roughly the
-    same number of atoms. Not as good as `smart_packing` but
-    faster and more scalable for-loop complexity of `O(batch_size)`.
-
-    Parameters:
-        num_nodes: List of the number of atoms per molecule for the entire global batch.
-            Must be of length `batch_size * ipu_batch_size`.
-
-        batch_size: The batch size per iteration, considering a single device and single
-            forward pass.
-            The global batch size is `batch_size * device_iterations * replication_factor * gradient_accumulation`
-
-    Returns:
-        packed_indices: A list of packs, each containing a list of indices, such that
-            if we collect `num_nodes` from the indices, then each pack has roughly the
-            same total number of atoms.
-    """
-    num_nodes = np.asarray(num_nodes)
-    argsort_num_nodes = np.argsort(num_nodes)
-    ipu_batch_size = int(len(num_nodes) / batch_size)
-
-    packed_indices = np.stack(
-        [
-            np.random.permutation(argsort_num_nodes[ii * ipu_batch_size : (ii + 1) * ipu_batch_size])
-            for ii in range(batch_size)
-        ],
-        axis=0,
-    ).T.tolist()
-    return packed_indices
-
-
-def hybrid_packing(num_nodes: List[int], batch_size: int) -> List[List[int]]:
-    """
-    Uses a combination of the `smart_packing` `O(n^2)` on the most important data points,
-    and the `fast_packing` `O(n)` on the average-sized data points.
-
-    Depending on the expected complexity
-
-    Parameters:
-        num_nodes: List of the number of atoms per molecule for the entire global batch.
-            Must be of length `batch_size * ipu_batch_size`.
-
-        batch_size: The batch size per iteration, considering a single device and single
-            forward pass.
-            The global batch size is `batch_size * device_iterations * replication_factor * gradient_accumulation`
-
-    Returns:
-        packed_indices: A list of packs, each containing a list of indices, such that
-            if we collect `num_nodes` from the indices, then each pack has roughly the
-            same total number of atoms.
-    """
-
-    # Determine the parameters based on the complexity of the smart-packing.
-    # The bigger the complexity, the more the `fast_packing` algorithm becomes
-    # statistically powerful, and the more speed benefits it provides.
-    smart_packing_complexity = len(num_nodes) ** 2 / batch_size
-    if smart_packing_complexity < 1e4:
-        return smart_packing(num_nodes=num_nodes, batch_size=batch_size)
-    elif smart_packing_complexity < 1e5:
-        big, small = 3, 6
-    else:
-        return fast_packing(num_nodes=num_nodes, batch_size=batch_size)
-
-    # Small datasets benefit from smart-packing, without compute burden
-    ipu_batch_size = int(len(num_nodes) / batch_size)
-    if len(num_nodes) < (big + small) * ipu_batch_size:
-        return smart_packing(num_nodes=num_nodes, batch_size=batch_size)
-
-    # Sort the list
-    num_nodes = np.asarray(num_nodes)
-    argsort_num_nodes = np.argsort(num_nodes)
-
-    # Smallest and biggest graphs are often outliers and will benefit from the `smart_packing`
-    biggest_graphs = argsort_num_nodes[-big * ipu_batch_size :]
-    smallest_graphs = argsort_num_nodes[: small * ipu_batch_size]
-    big_n_small_graphs = np.concatenate([biggest_graphs, smallest_graphs])
-    big_n_small_packs = smart_packing(num_nodes[big_n_small_graphs], batch_size=big + small)
-    big_n_small_indices = [big_n_small_graphs[pack] for pack in big_n_small_packs]
-    big_n_small_nodes = [num_nodes[pack] for pack in big_n_small_indices]
-
-    # Medium graphs will be packed faster
-    medium_graphs = argsort_num_nodes[small * ipu_batch_size : -big * ipu_batch_size]
-    medium_packs = fast_packing(num_nodes[medium_graphs], batch_size=batch_size - big - small)
-    medium_indices = [medium_graphs[pack] for pack in medium_packs]
-    medium_nodes = [num_nodes[pack] for pack in medium_indices]
-
-    # Pack the big/small with the medium in a smart way
-    big_n_small_sort = np.argsort(np.sum(np.stack(big_n_small_nodes, axis=1), axis=0))
-    medium_sort = np.argsort(np.sum(np.stack(medium_nodes, axis=1), axis=0))
-    packed_indices = [
-        np.concatenate([medium_indices[medium_sort[ii]], big_n_small_indices[big_n_small_sort[-ii]]])
-        for ii in range(len(medium_indices))
-    ]
-
-    return packed_indices
-
-
-def get_pack_sizes(packed_indices, num_nodes):
-    """
-    Get the number of atoms of each pack
-    """
-    pack_sums = []
-    for pack in packed_indices:
-        pack_sum = 0
-        for idx in pack:
-            pack_sum += num_nodes[idx]
-        pack_sums.append(pack_sum)
-    return pack_sums
-
-
-def estimate_max_pack_node_size(num_nodes: Iterable[int], batch_size: int, combined_batch_size: int):
-    """
-    Estimate the value of `max_num_nodes`, which represents the maximum number of nodes
-    needed in a batch to fit the data.
-
-    Parameters:
-        num_nodes: Number of nodes for all the graphs in the dataset
-        batch_size: The regular batch size per IPU
-        combined_batch_size: batch_size * device_iterations
-                             * replication_factor * gradient_accumulation
-
-    """
-
-    # Estimate the packing size needed
-    rand_indices = np.arange(len(num_nodes))
-    np.random.shuffle(rand_indices)
-    max_pack_size = 0
-    for ii in range(0, len(num_nodes), combined_batch_size):
-        this_indices = rand_indices[ii : ii + combined_batch_size]
-        choice = num_nodes[this_indices]
-        if len(choice) == combined_batch_size:
-            packed_indices = hybrid_packing(choice, batch_size)
-            max_pack_size = max(max_pack_size, max(get_pack_sizes(packed_indices, num_nodes[this_indices])))
-    max_pack_size_per_graph = max_pack_size / batch_size
-
-    return max_pack_size, max_pack_size_per_graph
-
-
-def node_to_pack_indices_mask(
-    packed_indices: Iterable[Iterable[int]], all_num_nodes: Iterable[int], max_pack_size: Optional[int] = None
-) -> Tuple[torch.Tensor, torch.Tensor]:
-    """
-    Given a list of packed indices, and the number of nodes in each graph,
-    return a tensor of shape (sum(all_num_nodes), 2) where the first column
-    is the pack index, and the second column is the node index within the pack.
-
-    Can be used to generate a dense packing of the nodes as follows:
-    ```
-    # node_features: A tensor of shape (num_nodes, num_node_features)
-    # num_packs: The number of packs desired
-    # max_nodes_per_pack: The maximum number of nodes per pack
-    # dense_pack: A tensor of shape (num_packs, max_nodes_per_pack, num_node_features)
-
-    dense_pack = torch.zeros([num_packs, max_nodes_per_pack, num_node_features])
-    dense_pack[pack_from_node_idx[:, 0], pack_from_node_idx[:, 1]] = node_features
-    ```
-
-    This is useful when using a Transformer, to avoid wasteful padding when the
-    the longest sequence is much longer than the average sequence length.
-
-    Parameters:
-        packed_indices: A list of lists of graph indices, where each sub-list
-                        represents a pack of graphs
-        all_num_nodes: The number of nodes in each graph
-        max_pack_size: The maximum number of nodes per pack. If None, will be
-                          infered from the provided packs.
-                          Useful to determine the shape of the `pack_attn_mask`.
-
-    Returns:
-        pack_from_node_idx: A tensor of shape (num_nodes, 2) where the first column
-                        is the pack index, and the second column is the node index within the pack.
-
-        pack_attn_mask: A tensor of shape (num_packs, max_pack_size, max_pack_size),
-                            that represents the attention masking for each pack,
-                            such that the graphs in the pack are masked out from each other.
-    """
-
-    all_num_nodes = torch.as_tensor(all_num_nodes, dtype=torch.long)
-    cumsum_num_nodes = torch.cumsum(all_num_nodes, dim=0)
-    if max_pack_size is None:
-        pack_sizes = get_pack_sizes(packed_indices, all_num_nodes)
-        max_pack_size = max(pack_sizes)
-
-    # Get the node indices associated to the packs, with 0 padding
-    pack_from_node_idx = torch.zeros(sum(all_num_nodes), 2, dtype=torch.long)
-    pack_attn_mask = []  # masks for the attention
-    for ii, pack in enumerate(packed_indices):
-        jj = 0  # Counter for the number of nodes in the pack
-        this_pack_attn_mask = torch.ones((max_pack_size, max_pack_size), dtype=torch.bool)
-        for graph_idx in pack:
-            num_nodes = all_num_nodes[graph_idx]
-            node_idx = torch.arange(cumsum_num_nodes[graph_idx] - num_nodes, cumsum_num_nodes[graph_idx])
-            this_pack_attn_mask[jj : jj + num_nodes, jj : jj + num_nodes] = False
-            pack_from_node_idx[node_idx, 0] = ii
-            pack_from_node_idx[node_idx, 1] = jj + torch.arange(num_nodes)
-            jj += num_nodes
-        pack_attn_mask.append(this_pack_attn_mask)
-    pack_attn_mask = torch.stack(pack_attn_mask, dim=0)
-
-    return pack_from_node_idx, pack_attn_mask
diff --git a/graphium/utils/spaces.py b/graphium/utils/spaces.py
index c7b6a7ac9..0ed17c0d7 100644
--- a/graphium/utils/spaces.py
+++ b/graphium/utils/spaces.py
@@ -22,7 +22,6 @@
 import graphium.nn.architectures as Architectures
 import graphium.utils.custom_lr as CustomLR
 import graphium.data.datamodule as Datamodules
-import graphium.ipu.ipu_losses as IPULosses
 import graphium.nn.pyg_layers as PygLayers
 import graphium.nn.residual_connections as Residuals
 import graphium.nn.encoders as Encoders
@@ -78,12 +77,6 @@
     "l1": torch.nn.L1Loss,
     "mae": torch.nn.L1Loss,
     "hybrid_ce": Losses.HybridCELoss,
-    "bce_ipu": IPULosses.BCELossIPU,
-    "bce_logits_ipu": IPULosses.BCEWithLogitsLossIPU,
-    "mse_ipu": IPULosses.MSELossIPU,
-    "mae_ipu": IPULosses.L1LossIPU,
-    "l1_ipu": IPULosses.L1LossIPU,
-    "hybrid_ce_ipu": IPULosses.HybridCELossIPU,
 }
 
 
diff --git a/install_ipu.sh b/install_ipu.sh
deleted file mode 100755
index a21022bdb..000000000
--- a/install_ipu.sh
+++ /dev/null
@@ -1,112 +0,0 @@
-"""
---------------------------------------------------------------------------------
-Copyright (c) 2023 Graphcore Limited.
-Use of this software is subject to the terms and conditions outlined in the LICENSE file.
-Unauthorized modification, distribution, or use is prohibited. Provided 'as is' without
-warranties of any kind.
-
-Graphcore Limited is not liable
-for any damages arising from its use.
-Refer to the LICENSE file for the full terms and conditions.
---------------------------------------------------------------------------------
-"""
-
-
-#!/bin/bash
-
-# Default location for the virtual environment
-default_venv_name=".graphium_ipu"
-
-# Allow the user to specify the location of their virtual environment
-# If not specified, use the default location
-venv_name=${1:-$default_venv_name}
-
-# Constants
-sdk_compressed_file="poplar_sdk-ubuntu_20_04-3.3.0-208993bbb7.tar.gz"
-sdk_wheel_file="poptorch-3.3.0+113432_960e9c294b_ubuntu_20_04-cp38-cp38-linux_x86_64.whl"
-sdk_url="https://downloads.graphcore.ai/direct?package=poplar-poplar_sdk_ubuntu_20_04_3.3.0_208993bbb7-3.3.0&file=${sdk_compressed_file}"
-sdk_path="${venv_name}/poplar_sdk-ubuntu_20_04-3.3.0+1403-208993bbb7"
-
-# Check for Python3 and pip
-if ! command -v python3 &>/dev/null; then
-    echo "Python3 is required but it's not installed. Exiting."
-    exit 1
-fi
-
-if ! command -v pip3 &>/dev/null; then
-    echo "pip3 is required but it's not installed. Exiting."
-    exit 1
-fi
-
-# Remove existing venv directory if it exists
-if [[ -d $venv_name ]]; then
-    echo "Removing existing virtual environment directory..."
-    rm -rf $venv_name
-fi
-
-# Create the virtual environment
-echo "Creating virtual environment..."
-mkdir -p $venv_name
-python3 -m venv $venv_name
-source $venv_name/bin/activate
-
-# Update pip to the latest version
-echo "Upgrading pip..."
-python3 -m pip install --upgrade pip
-
-# Download the Poplar SDK
-echo "Downloading Poplar SDK..."
-wget -q -O "${venv_name}/${sdk_compressed_file}" "$sdk_url"
-
-# Check the wget exit status
-if [ $? -ne 0 ]; then
-    echo "Failed to download Poplar SDK. Exiting."
-    exit 1
-fi
-
-# Unzip the SDK file
-echo "Extracting Poplar SDK..."
-tar -xzf "$venv_name/$sdk_compressed_file" -C $venv_name
-
-# Install the PopTorch wheel
-echo "Installing PopTorch..."
-python3 -m pip install "${sdk_path}/${sdk_wheel_file}"
-
-# Enable Poplar SDK (including Poplar and PopART)
-echo "Enabling Poplar SDK..."
-source ${sdk_path}/enable
-
-# Install the IPU specific and Graphium requirements
-echo "Installing IPU specific and Graphium requirements..."
-python3 -m pip install -r requirements_ipu.txt
-
-# Install Graphium in dev mode
-echo "Installing Graphium in dev mode..."
-python3 -m pip install --no-deps -e .
-
-# This is a quick test make sure poptorch is correctly installed
-if python3 -c "import poptorch;print('poptorch installed correctly')" &> /dev/null; then
-    echo "Installation completed successfully."
-else
-    echo "Installation was not successful. Please check the logs and try again."
-    exit 1  # Exit with status code 1 to indicate failure
-fi
-
-# Download the datafiles (Total ~ 10Mb - nothing compared to the libraries)
-echo "Downloading the sub-datasets consisting on the ToyMix dataset"
-toymix_dir=expts/data/neurips2023/small-dataset/
-mkdir -p $toymix_dir
-
-base_url="https://storage.googleapis.com/graphium-public/datasets/neurips_2023/Small-dataset/"
-files=("ZINC12k.csv.gz" "Tox21-7k-12-labels.csv.gz" "qm9.csv.gz" "qm9_random_splits.pt" "Tox21_random_splits.pt" "ZINC12k_random_splits.pt")
-
-for file in "${files[@]}"; do
-    if [ ! -f "${toymix_dir}${file}" ]; then
-        echo "Downloading ${file}..."
-        wget -P "${toymix_dir}" "${base_url}${file}"
-    else
-        echo "${file} already exists. Skipping..."
-    fi
-done
-
-echo "Data has been successfully downloaded."
\ No newline at end of file
diff --git a/mkdocs.yml b/mkdocs.yml
index 0d9f34bfd..0fb199fa6 100644
--- a/mkdocs.yml
+++ b/mkdocs.yml
@@ -23,7 +23,6 @@ nav:
       - graphium.data: api/graphium.data.md
       - graphium.utils: api/graphium.utils.md
       - graphium.config: api/graphium.config.md
-      - graphium.ipu: api/graphium.ipu.md
       - graphium.finetuning: api/graphium.finetuning.md
   - Tutorials:
       - feature_processing:
@@ -36,7 +35,6 @@ nav:
           - Using GNN layers: tutorials/gnn/using_gnn_layers.ipynb
       - model_training:
           - Simple Molecular Model: tutorials/model_training/simple-molecular-model.ipynb
-          - Training on IPU: tutorials/model_training/running-multitask-ipu.ipynb
   - Design: design.md
   - Datasets: datasets.md
   - Pretrained Models: pretrained_models.md
diff --git a/pyproject.toml b/pyproject.toml
index 78a5869da..364d8fd1b 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -104,7 +104,6 @@ filterwarnings = [
     "ignore::DeprecationWarning:pkg_resources.*:",
 ]
 markers = [
-    "ipu: marks tests that are specific to the IPU (deselect with '-m \"not ipu\"')",
 ]
 
 [tool.coverage.run]
diff --git a/scripts/ipu_start.sh b/scripts/ipu_start.sh
deleted file mode 100644
index b50ffcbd9..000000000
--- a/scripts/ipu_start.sh
+++ /dev/null
@@ -1,25 +0,0 @@
-"""
---------------------------------------------------------------------------------
-Copyright (c) 2023 Valence Labs, Recursion Pharmaceuticals and Graphcore Limited.
-
-Use of this software is subject to the terms and conditions outlined in the LICENSE file.
-Unauthorized modification, distribution, or use is prohibited. Provided 'as is' without
-warranties of any kind.
-
-Valence Labs, Recursion Pharmaceuticals and Graphcore Limited are not liable for any damages arising from its use.
-Refer to the LICENSE file for the full terms and conditions.
---------------------------------------------------------------------------------
-"""
-
-
-"""
-Start the ipu environment and SDK
-"""
-
-source /opt/gc/sdk-3.0.0+1128/poplar-ubuntu_20_04-3.0.0+5468-0379b9a65d/enable.sh
-source /opt/gc/sdk-3.0.0+1128/popart-ubuntu_20_04-3.0.0+5468-0379b9a65d/enable.sh
-
-source ~/.venv/graphium_ipu/bin/activate # Change to your path
-
-export VISUAL=vim
-export EDITOR="$VISUAL"
diff --git a/scripts/ipu_venv.sh b/scripts/ipu_venv.sh
deleted file mode 100644
index 826fcfa12..000000000
--- a/scripts/ipu_venv.sh
+++ /dev/null
@@ -1,30 +0,0 @@
-"""
---------------------------------------------------------------------------------
-Copyright (c) 2023 Valence Labs, Recursion Pharmaceuticals and Graphcore Limited.
-
-Use of this software is subject to the terms and conditions outlined in the LICENSE file.
-Unauthorized modification, distribution, or use is prohibited. Provided 'as is' without
-warranties of any kind.
-
-Valence Labs, Recursion Pharmaceuticals and Graphcore Limited are not liable for any damages arising from its use.
-Refer to the LICENSE file for the full terms and conditions.
---------------------------------------------------------------------------------
-"""
-
-
-"""
-Create the pip environment for IPU
-"""
-
-## Uncomment this to create the folder for the environment
-# mkdir ~/.venv                              # Create the folder for the environment
-# python3 -m venv ~/.venv/graphium_ipu           # Create the environment
-# source ~/.venv/graphium_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
-pip install /opt/gc/sdk-3.0.0+1128/poptorch-3.0.0+84519_672c9cbc7f_ubuntu_20_04-cp38-cp38-linux_x86_64.whl
-pip install -r requirements.txt
-pip install -e .


Test patch:
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index af85c8616..f84debd95 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -16,8 +16,13 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: ["3.09", "3.10", "3.11"]
-        pytorch-version: ["2.0"]
+        include:
+          - python-version: "3.10"
+            pytorch-version: "2.0"
+          - python-version: "3.11"
+            pytorch-version: "2.0"
+          - python-version: "3.12"
+            pytorch-version: "2.3"
 
     runs-on: "ubuntu-latest"
     timeout-minutes: 30
@@ -55,9 +60,6 @@ jobs:
       - name: Install C++ library
         run: cd graphium/graphium_cpp && git clone https://github.com/pybind/pybind11.git && export PYTHONPATH=$PYTHONPATH:./pybind11 && python -m pip install . && cd ../..
 
-      - name: Run tests
-        run: pytest -m 'not ipu'
-
       - name: Test CLI
         run: graphium --help
 
diff --git a/.github/workflows/test_ipu.yml b/.github/workflows/test_ipu.yml
deleted file mode 100644
index 886c4c2b7..000000000
--- a/.github/workflows/test_ipu.yml
+++ /dev/null
@@ -1,69 +0,0 @@
-name: test-ipu
-
-on:
-  push:
-    branches: ["main"]
-    tags: ["*"]
-  pull_request:
-    branches:
-      - "*"
-      - "!gh-pages"
-  schedule:
-    - cron: "0 4 * * *"
-
-jobs:
-  test-ipu:
-    strategy:
-      fail-fast: false
-      matrix:
-        python-version: ["3.8"]
-        pytorch-version: ["2.0"]
-
-    runs-on: "ubuntu-20.04"
-    timeout-minutes: 30
-
-    defaults:
-      run:
-        shell: bash -l {0}
-
-    name: |
-        poptorch_env - 
-        python=${{ matrix.python-version }} -
-        pytorch=${{ matrix.pytorch-version }}
-
-    steps:
-      - name: Checkout the code
-        uses: actions/checkout@v3
-
-      - name: Activate SDK + Install Requirements
-        run: |
-          python3 -m pip install --upgrade pip
-          wget -q -O 'poplar_sdk-ubuntu_20_04-3.3.0-208993bbb7.tar.gz' 'https://downloads.graphcore.ai/direct?package=poplar-poplar_sdk_ubuntu_20_04_3.3.0_208993bbb7-3.3.0&file=poplar_sdk-ubuntu_20_04-3.3.0-208993bbb7.tar.gz'
-          tar -xzf poplar_sdk-ubuntu_20_04-3.3.0-208993bbb7.tar.gz
-          python3 -m pip install poplar_sdk-ubuntu_20_04-3.3.0+1403-208993bbb7/poptorch-3.3.0+113432_960e9c294b_ubuntu_20_04-cp38-cp38-linux_x86_64.whl
-          # Enable Poplar SDK (including Poplar and PopART)
-          source poplar_sdk-ubuntu_20_04-3.3.0+1403-208993bbb7/enable 
-          
-          python -c "import poptorch"
-
-          # Download the datafiles (Total ~ 10Mb - nothing compared to the libraries)
-          wget https://storage.googleapis.com/graphium-public/datasets/neurips_2023/Small-dataset/ZINC12k.csv.gz
-          wget https://storage.googleapis.com/graphium-public/datasets/neurips_2023/Small-dataset/Tox21-7k-12-labels.csv.gz
-          wget https://storage.googleapis.com/graphium-public/datasets/neurips_2023/Small-dataset/qm9.csv.gz
-
-
-          # Install the IPU specific and graphium requirements
-          pip install -r requirements_ipu.txt
-          # Install Graphium in dev mode
-          python -m pip install --no-deps -e .
-          python3 -m pytest -m 'not skip_ipu'
-
-      - name: Codecov Upload
-        uses: codecov/codecov-action@v3
-        with:
-          files: ./coverage.xml
-          flags: unittests
-          name: codecov-umbrella
-          fail_ci_if_error: false
-          verbose: false
-          env_vars: ${{ matrix.python-version }},${{ matrix.pytorch-version }}
diff --git a/graphium/cli/train_finetune_test.py b/graphium/cli/train_finetune_test.py
index 4cea1ee8c..29152bdb5 100644
--- a/graphium/cli/train_finetune_test.py
+++ b/graphium/cli/train_finetune_test.py
@@ -43,7 +43,6 @@
 
 TESTING_ONLY_CONFIG_KEY = "testing_only"
 
-
 @hydra.main(version_base=None, config_path="../../expts/hydra-configs", config_name="main")
 def cli(cfg: DictConfig) -> None:
     """
@@ -52,41 +51,6 @@ def cli(cfg: DictConfig) -> None:
     return run_training_finetuning_testing(cfg)
 
 
-def get_training_batch_size(cfg):
-    """
-    WARNING: This MUST be called after accelerator overrides have been applied
-    (i.e. after `load_accelerator` has been called)
-    """
-    try:
-        # Navigate through the nested dictionaries and get the training batch size
-        batch_size_training = cfg.get("datamodule", {}).get("args", {}).get("batch_size_training", 1)
-
-        # Ensure that the extracted value is an integer
-        return int(batch_size_training)
-    except Exception as e:
-        print(f"An error occurred: {e}")
-
-    # Return default value if an error occurred
-    return 1
-
-
-def get_training_device_iterations(cfg):
-    try:
-        ipu_config = cfg.get("accelerator", {}).get("ipu_config", [])
-        for item in ipu_config:
-            if "deviceIterations" in item:
-                # Extract the number between parentheses
-                start = item.find("(") + 1
-                end = item.find(")")
-                if start != 0 and end != -1:
-                    return int(item[start:end])
-    except Exception as e:
-        print(f"An error occurred: {e}")
-
-    # Return default value if deviceIterations is not found or an error occurred
-    return 1
-
-
 def run_training_finetuning_testing(cfg: DictConfig) -> None:
     """
     The main (pre-)training and fine-tuning loop.
diff --git a/tests/config_test_ipu_dataloader.yaml b/tests/config_test_dataloader.yaml
similarity index 93%
rename from tests/config_test_ipu_dataloader.yaml
rename to tests/config_test_dataloader.yaml
index 3f63bfd3d..c1ef946d1 100644
--- a/tests/config_test_ipu_dataloader.yaml
+++ b/tests/config_test_dataloader.yaml
@@ -1,22 +1,14 @@
-# Testing the multitask pipeline with the QM9 dataset on IPU, by splitting it up into three tasks: homo, alpha and cv.
+# Testing the multitask pipeline with the QM9 dataset, by splitting it up into three tasks: homo, alpha and cv.
 constants:
-  name: &name test_ipu #qm9_full
+  name: &name test_dataloader
   seed: &seed 42
   raise_train_error: true   # Whether the code should raise an error if it crashes during training
 
 accelerator:
-  type: ipu  # cpu or ipu or gpu
+  type: cpu  # cpu or gpu
   config_override:
     datamodule:
       args:
-        ipu_dataloader_training_opts:
-          mode: async
-          max_num_nodes_per_graph: 20 # train max nodes: 20, max_edges: 54
-          max_num_edges_per_graph: 60
-        ipu_dataloader_inference_opts:
-          mode: async
-          max_num_nodes_per_graph: 16 # valid max nodes: 51, max_edges: 118
-          max_num_edges_per_graph: 120
         # Data handling-related
         batch_size_training: 6
         batch_size_inference: 6
@@ -25,9 +17,6 @@ accelerator:
         precision: 16
         accumulate_grad_batches: 4
 
-  ipu_config:
-    - deviceIterations(2)
-
 datamodule:
   module_type: "MultitaskFromSmilesDataModule"
   args: # Matches that in the test_multitask_datamodule.py case.
@@ -61,6 +50,9 @@ datamodule:
         weights_type: null                # This may not always be provided
         task_level: graph
     # Featurization
+    prepare_dict_or_graph: pyg:graph
+    featurization_n_jobs: 0
+    featurization_progress: True
     featurization:
       atom_property_list_onehot: [atomic-number, valence]
       atom_property_list_float: [mass, electronegativity, in-ring]
@@ -266,8 +258,8 @@ predictor:
     homo: ["mae"]
     alpha: ["mae"]
   loss_fun:
-    homo: mse_ipu
-    alpha: mse_ipu
+    homo: mse
+    alpha: mse
   random_seed: *seed
   optim_kwargs:
     lr: 1.e-3
@@ -297,4 +289,4 @@ trainer:
     every_n_epochs: 1
   trainer:
     max_epochs: 2
-    min_epochs: 1
+    min_epochs: 1
\ No newline at end of file
diff --git a/tests/config_test_ipu_dataloader_multitask.yaml b/tests/config_test_ipu_dataloader_multitask.yaml
deleted file mode 100644
index 563222d8d..000000000
--- a/tests/config_test_ipu_dataloader_multitask.yaml
+++ /dev/null
@@ -1,336 +0,0 @@
-# Testing the gcn model with the PCQMv2 dataset on IPU.
-constants:
-  name: &name neurips2023_small_data_gcn
-  seed: &seed 42
-  raise_train_error: true   # Whether the code should raise an error if it crashes during training
-
-accelerator:
-  type: ipu  # cpu or ipu or gpu
-  config_override:
-    datamodule:
-      args:
-        ipu_dataloader_training_opts:
-          mode: async
-          max_num_nodes_per_graph: 44 # train max nodes: 20, max_edges: 54
-          max_num_edges_per_graph: 80
-        ipu_dataloader_inference_opts:
-          mode: async
-          max_num_nodes_per_graph: 44 # valid max nodes: 51, max_edges: 118
-          max_num_edges_per_graph: 80
-        # Data handling-related
-        batch_size_training: 50
-        batch_size_inference: 50
-    predictor:
-      optim_kwargs:
-        loss_scaling: 1024
-    trainer:
-      trainer:
-        precision: 16
-        accumulate_grad_batches: 4
-
-  ipu_config:
-    - deviceIterations(5) # IPU would require large batches to be ready for the model.
-    - replicationFactor(1)
-    # - enableProfiling("graph_analyser")       # The folder where the profile will be stored
-    # - enableExecutableCaching("pop_compiler_cache")
-    - TensorLocations.numIOTiles(128)
-    - _Popart.set("defaultBufferingDepth", 128)
-    - Precision.enableStochasticRounding(True)
-    - useIpuModel(True)
-
-# accelerator:
-#   type: cpu  # cpu or ipu or gpu
-#   config_override:
-#     datamodule:
-#       batch_size_training: 64
-#       batch_size_inference: 256
-#     trainer:
-#       trainer:
-#         precision: 32
-#         accumulate_grad_batches: 1
-
-datamodule:
-  module_type: "MultitaskFromSmilesDataModule"
-  args: # Matches that in the test_multitask_datamodule.py case.
-    task_specific_args:   # To be replaced by a new class "DatasetParams"
-      qm9:
-        df: null
-        df_path: qm9.csv.gz
-        # wget https://storage.googleapis.com/graphium-public/datasets/neurips_2023/Small-dataset/qm9.csv.gz
-        # or set path as the URL directly
-        smiles_col: "smiles"
-        label_cols: ["A", "B", "C", "mu", "alpha", "homo", "lumo", "gap", "r2", "zpve", "u0", "u298", "h298", "g298", "cv", "u0_atom", "u298_atom", "h298_atom", "g298_atom"]
-        sample_size: 2000 # use sample_size for test
-        seed: *seed
-        task_level: graph
-        label_normalization:
-          normalize_val_test: True
-          method: "normal"
-
-      tox21:
-        df: null
-        df_path: Tox21-7k-12-labels.csv.gz
-        # wget https://storage.googleapis.com/graphium-public/datasets/neurips_2023/Small-dataset/Tox21-7k-12-labels.csv.gz
-        # or set path as the URL directly
-        smiles_col: "smiles"
-        label_cols: ["NR-AR", "NR-AR-LBD", "NR-AhR", "NR-Aromatase", "NR-ER", "NR-ER-LBD", "NR-PPAR-gamma", "SR-ARE", "SR-ATAD5", "SR-HSE", "SR-MMP", "SR-p53"]
-        sample_size: 2000 # use sample_size for test
-        seed: *seed
-        task_level: graph
-
-      zinc:
-        df: null
-        df_path: ZINC12k.csv.gz
-        # df_path: data/neurips2023/small-dataset/ZINC12k.csv.gz
-        # wget https://storage.googleapis.com/graphium-public/datasets/neurips_2023/Small-dataset/ZINC12k.csv.gz
-        # or set path as the URL directly
-        smiles_col: "smiles"
-        label_cols: ["SA", "logp", "score"]
-        sample_size: 2000 # use sample_size for test
-        seed: *seed
-        task_level: graph
-        label_normalization:
-          normalize_val_test: True
-          method: "normal"
-
-    # Featurization
-    # processed_graph_data_path: "../datacache/neurips2023-small/"
-    featurization:
-    # OGB: ['atomic_num', 'degree', 'possible_formal_charge', 'possible_numH' (total-valence),
-    # 'possible_number_radical_e', 'possible_is_aromatic', 'possible_is_in_ring',
-    # 'num_chiral_centers (not included yet)']
-      atom_property_list_onehot: [atomic-number, group, period, total-valence]
-      atom_property_list_float: [degree, formal-charge, radical-electron, aromatic, in-ring]
-      # OGB: ['possible_bond_type', 'possible_bond_stereo', 'possible_is_in_ring']
-      edge_property_list: [bond-type-onehot, stereo, in-ring]
-      add_self_loop: False
-      explicit_H: False # if H is included
-      use_bonds_weights: False
-      pos_encoding_as_features: # encoder dropout 0.18
-        pos_types:
-          lap_eigvec:
-            pos_level: node
-            pos_type: laplacian_eigvec
-            num_pos: 8
-            normalization: "none" # nomrlization already applied on the eigen vectors
-            disconnected_comp: True # if eigen values/vector for disconnected graph are included
-          lap_eigval:
-            pos_level: node
-            pos_type: laplacian_eigval
-            num_pos: 8
-            normalization: "none" # nomrlization already applied on the eigen vectors
-            disconnected_comp: True # if eigen values/vector for disconnected graph are included
-          rw_pos: # use same name as pe_encoder
-            pos_level: node
-            pos_type: rw_return_probs
-            ksteps: 16
-
-    num_workers: -1 # -1 to use all
-    persistent_workers: False # if use persistent worker at the start of each epoch.
-    # Using persistent_workers false might make the start of each epoch very long.
-
-
-architecture:
-  model_type: FullGraphMultiTaskNetwork
-  mup_base_path: null
-  pre_nn:   # Set as null to avoid a pre-nn network
-    out_dim: 16
-    hidden_dims: 16
-    depth: 1
-    activation: relu
-    last_activation: none
-    dropout: &dropout 0.1
-    normalization: &normalization layer_norm
-    last_normalization: *normalization
-    residual_type: none
-
-  pre_nn_edges:  null # Set as null to avoid a pre-nn network
-
-  pe_encoders:
-    out_dim: 32
-    pool: "sum" #"mean" "max"
-    last_norm: None #"batch_norm", "layer_norm"
-    encoders: #la_pos |  rw_pos
-      la_pos:  # Set as null to avoid a pre-nn network
-        encoder_type: "laplacian_pe"
-        input_keys: ["laplacian_eigvec", "laplacian_eigval"]
-        output_keys: ["feat"]
-        hidden_dim: 64
-        out_dim: 32
-        model_type: 'DeepSet' #'Transformer' or 'DeepSet'
-        num_layers: 2
-        num_layers_post: 1 # Num. layers to apply after pooling
-        dropout: 0.1
-        first_normalization: "none" #"batch_norm" or "layer_norm"
-      rw_pos:
-        encoder_type: "mlp"
-        input_keys: ["rw_return_probs"]
-        output_keys: ["feat"]
-        hidden_dim: 64
-        out_dim: 32
-        num_layers: 2
-        dropout: 0.1
-        normalization: "layer_norm" #"batch_norm" or "layer_norm"
-        first_normalization: "layer_norm" #"batch_norm" or "layer_norm"
-
-
-
-  gnn:  # Set as null to avoid a post-nn network
-    in_dim: 16 # or otherwise the correct value
-    out_dim: &gnn_dim 16
-    hidden_dims: *gnn_dim
-    depth: 1
-    activation: gelu
-    last_activation: none
-    dropout: 0.1
-    normalization: "layer_norm"
-    last_normalization: *normalization
-    residual_type: simple
-    virtual_node: 'none'
-    layer_type: 'pyg:gcn' #pyg:gine #'pyg:gps' # pyg:gated-gcn, pyg:gine,pyg:gps
-    layer_kwargs: null # Parameters for the model itself. You could define dropout_attn: 0.1
-
-
-  graph_output_nn:
-    graph:
-      pooling: [sum]
-      out_dim: *gnn_dim
-      hidden_dims: *gnn_dim
-      depth: 1
-      activation: relu
-      last_activation: none
-      dropout: *dropout
-      normalization: *normalization
-      last_normalization: "none"
-      residual_type: none
-
-  task_heads:
-    qm9:
-      task_level: graph
-      out_dim: 19
-      hidden_dims: 16
-      depth: 1
-      activation: relu
-      last_activation: none
-      dropout: *dropout
-      normalization: *normalization
-      last_normalization: "none"
-      residual_type: none
-    tox21:
-      task_level: graph
-      out_dim: 12
-      hidden_dims: 16
-      depth: 1
-      activation: relu
-      last_activation: sigmoid
-      dropout: *dropout
-      normalization: *normalization
-      last_normalization: "none"
-      residual_type: none
-    zinc:
-      task_level: graph
-      out_dim: 3
-      hidden_dims: 16
-      depth: 2
-      activation: relu
-      last_activation: none
-      dropout: *dropout
-      normalization: *normalization
-      last_normalization: "none"
-      residual_type: none
-
-#Task-specific
-predictor:
-  metrics_on_progress_bar:
-    qm9: ["mae"]
-    tox21: ["auroc"]
-    zinc: ["mae"]
-  loss_fun:
-    qm9: mae_ipu
-    tox21: bce_ipu
-    zinc: mae_ipu
-  random_seed: *seed
-  optim_kwargs:
-    lr: 4.e-5 # warmup can be scheduled using torch_scheduler_kwargs
-    # weight_decay: 1.e-7
-  torch_scheduler_kwargs:
-    module_type: WarmUpLinearLR
-    max_num_epochs: &max_epochs 1
-    warmup_epochs: 1
-    verbose: False
-  scheduler_kwargs:
-  #  monitor: &monitor qm9/mae/train
-  #  mode: min
-  #  frequency: 1
-  target_nan_mask: null # null: no mask, 0: 0 mask, ignore-flatten, ignore-mean-per-label
-  multitask_handling: flatten # flatten, mean-per-label
-
-# Task-specific
-metrics:
-  qm9: &qm9_metrics
-    - name: mae
-      metric: mae_ipu
-      target_nan_mask: null
-      multitask_handling: flatten
-      threshold_kwargs: null
-    - name: pearsonr
-      metric: pearsonr_ipu
-      threshold_kwargs: null
-      target_nan_mask: null
-      multitask_handling: mean-per-label
-    - name: r2_score
-      metric: r2_score_ipu
-      target_nan_mask: null
-      multitask_handling: mean-per-label
-      threshold_kwargs: null
-  tox21:
-    - name: auroc
-      metric: auroc_ipu
-      task: binary
-      multitask_handling: mean-per-label
-      threshold_kwargs: null
-    - name: avpr
-      metric: average_precision_ipu
-      task: binary
-      multitask_handling: mean-per-label
-      threshold_kwargs: null
-    - name: f1 > 0.5
-      metric: f1
-      multitask_handling: mean-per-label
-      target_to_int: True
-      num_classes: 2
-      average: micro
-      threshold_kwargs: &threshold_05
-        operator: greater
-        threshold: 0.5
-        th_on_preds: True
-        th_on_target: True
-    - name: precision > 0.5
-      metric: precision
-      multitask_handling: mean-per-label
-      average: micro
-      threshold_kwargs: *threshold_05
-  zinc: *qm9_metrics
-
-trainer:
-  seed: *seed
-  logger:
-    save_dir: logs/neurips2023-small/
-    name: *name
-    project: *name
-  #early_stopping:
-  #  monitor: *monitor
-  #  min_delta: 0
-  #  patience: 10
-  #  mode: &mode min
-  model_checkpoint:
-    dirpath: models_checkpoints/neurips2023-small-gcn/
-    filename: *name
-    # monitor: *monitor
-    # mode: *mode
-    # save_top_k: 1
-    save_last: True
-  trainer:
-    max_epochs: *max_epochs
-    min_epochs: 1
-    check_val_every_n_epoch: 20
diff --git a/tests/dummy-pretrained-model.ckpt b/tests/dummy-pretrained-model.ckpt
index e3df0fa4c..41e4df2b4 100644
Binary files a/tests/dummy-pretrained-model.ckpt and b/tests/dummy-pretrained-model.ckpt differ
diff --git a/tests/test_attention.py b/tests/test_attention.py
index 28b9cd2a1..bedea4933 100644
--- a/tests/test_attention.py
+++ b/tests/test_attention.py
@@ -21,8 +21,8 @@
 import torch
 import unittest as ut
 from torch_geometric.data import Data, Batch
+from torch_geometric.utils import to_dense_batch
 from copy import deepcopy
-from graphium.ipu.to_dense_batch import to_dense_batch
 
 from graphium.nn.base_layers import MultiheadAttentionMup
 
@@ -65,12 +65,11 @@ def test_attention_class(self):
         attention_layer_bias = MultiheadAttentionMup(biased_attention=True, **self.attn_kwargs)
         attention_layer_bias.eval()
 
-        h_dense, mask, _ = to_dense_batch(
+        h_dense, mask = to_dense_batch(
             bg.feat,
             batch=bg.batch,
             batch_size=None,
-            max_num_nodes_per_graph=None,
-            drop_nodes_last_graph=False,
+            max_num_nodes=None,
         )
         # attn_bias [batch, num_heads, nodes, nodes]
         nodes = h_dense.size()[1]
diff --git a/tests/test_base_layers.py b/tests/test_base_layers.py
index 2093619f2..6e153d480 100644
--- a/tests/test_base_layers.py
+++ b/tests/test_base_layers.py
@@ -19,10 +19,10 @@
 import torch
 import unittest as ut
 from torch_geometric.data import Data, Batch
+from torch_geometric.utils import to_dense_batch, dense_to_sparse
 from copy import deepcopy
 
 from graphium.nn.base_layers import DropPath, TransformerEncoderLayerMup
-from graphium.ipu.to_dense_batch import to_dense_batch, to_sparse_batch
 
 
 class test_Base_Layers(ut.TestCase):
@@ -78,18 +78,19 @@ def test_transformer_encoder_layer_mup(self):
             biased_attention=False, d_model=self.in_dim, nhead=1, dim_feedforward=4 * self.in_dim
         )
 
-        feat_dense, key_padding_mask, idx = to_dense_batch(
+        feat_dense, key_padding_mask = to_dense_batch(
             feat_in,
             batch=bg.batch,
             batch_size=self.batch_size,
-            max_num_nodes_per_graph=self.max_num_nodes_per_graph,
-            drop_nodes_last_graph=False,
+            max_num_nodes=self.max_num_nodes_per_graph,
         )
-        attn_mask = None
-        key_padding_mask = ~key_padding_mask
 
+        key_padding_mask = ~key_padding_mask
         h_out_dense = layer.forward(feat_dense)
-
-        h_out = to_sparse_batch(h_out_dense, mask_idx=idx)
+        h_out = h_out_dense[~key_padding_mask]
 
         self.assertEqual(h_out.shape, feat_in.shape)
+
+
+if __name__ == "__main__":
+    ut.main()
diff --git a/tests/test_ipu_dataloader.py b/tests/test_ipu_dataloader.py
deleted file mode 100644
index 2b018206e..000000000
--- a/tests/test_ipu_dataloader.py
+++ /dev/null
@@ -1,256 +0,0 @@
-"""
---------------------------------------------------------------------------------
-Copyright (c) 2023 Valence Labs, Recursion Pharmaceuticals and Graphcore Limited.
-
-Use of this software is subject to the terms and conditions outlined in the LICENSE file.
-Unauthorized modification, distribution, or use is prohibited. Provided 'as is' without
-warranties of any kind.
-
-Valence Labs, Recursion Pharmaceuticals and Graphcore Limited are not liable for any damages arising from its use.
-Refer to the LICENSE file for the full terms and conditions.
---------------------------------------------------------------------------------
-"""
-
-
-# General imports
-import yaml
-import unittest as ut
-import numpy as np
-from copy import deepcopy
-from warnings import warn
-from unittest.mock import patch
-from lightning import Trainer, LightningModule
-from functools import partial
-import pytest
-from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, Union
-
-import torch
-from torch.utils.data.dataloader import default_collate
-
-
-def random_packing(num_nodes, batch_size):
-    ipu_batch_size = int(len(num_nodes) / batch_size)
-    indices = np.arange(len(num_nodes))
-    np.random.shuffle(indices)
-    indices = np.reshape(indices, (ipu_batch_size, batch_size)).tolist()
-    return indices
-
-
-def global_batch_collator(batch_size, batches):
-    packs = []
-    for pack_idx in range(0, len(batches), batch_size):
-        packs.append(default_collate(batches[pack_idx : pack_idx + batch_size]))
-    global_batch = default_collate(packs)
-    global_batch = (global_batch[0], tuple(global_batch[1]))
-    return global_batch
-
-
-@pytest.mark.ipu
-class test_DataLoading(ut.TestCase):
-    class SimpleLightningTesting(LightningModule):
-        # Create a basic Ligthning for testing the batch sizes
-        def __init__(self, batch_size, node_feat_size, edge_feat_size, num_batch) -> None:
-            super().__init__()
-            self.batch_size = batch_size
-            self.node_feat_size = node_feat_size
-            self.edge_feat_size = edge_feat_size
-            self.layer = torch.nn.Linear(node_feat_size, 1)
-            self.loss_fn = torch.nn.L1Loss()
-            self.num_batch = num_batch
-
-        def validation_step(self, batch, batch_idx):
-            self.assert_shapes(batch, batch_idx, "val")
-            loss = self.forward(batch)
-            return loss
-
-        def training_step(self, batch, batch_idx):
-            self.assert_shapes(batch, batch_idx, "train")
-            loss = self.forward(batch)
-            return loss
-
-        def forward(self, batch):
-            out = self.layer(batch[1][0]).squeeze(-1)
-            loss = self.loss_fn(out, batch[0])
-            return loss
-
-        def assert_shapes(self, batch, batch_idx, 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}"
-
-            # 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}"
-
-            # 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}"
-
-        def configure_optimizers(self):
-            return torch.optim.Adam(self.parameters(), lr=1e-3)
-
-    class DatasetTesting(torch.utils.data.Dataset):
-        # Create a simple dataset for testing the Lightning integration
-        def __init__(self, labels, node_features, edge_features):
-            self.labels = labels
-            self.node_features = node_features
-            self.edge_features = edge_features
-
-        def __len__(self):
-            return len(self.labels)
-
-        def __getitem__(self, idx):
-            # [label, [feat1, feat2]]
-            return [self.labels[idx], [self.node_features[idx], self.edge_features[idx]]]
-
-    # @pytest.mark.skip
-    def test_poptorch_simple_deviceiterations_gradient_accumulation(self):
-        """
-        Test a simple version of the device-iterations and gradient accumulation
-        to make sure that the dataloader and models handle them correcly.
-        """
-        from lightning_graphcore import IPUStrategy
-
-
-        with patch("poptorch.ipuHardwareIsAvailable", return_value=True):
-            with patch("lightning_graphcore.accelerator._IPU_AVAILABLE", new=True):
-                import poptorch
-
-                assert poptorch.ipuHardwareIsAvailable()
-                from lightning_graphcore.accelerator import _IPU_AVAILABLE
-
-                assert _IPU_AVAILABLE is True
-
-                # Initialize constants
-                gradient_accumulation = 2
-                device_iterations = 3
-                batch_size = 5
-                num_replicate = 7
-                node_feat_size = 11
-                edge_feat_size = 13
-
-                # Initialize the batch info and poptorch options
-                opts = poptorch.Options()
-                opts.useIpuModel(True)
-                opts.deviceIterations(device_iterations)
-                training_opts = deepcopy(opts)
-                training_opts.Training.gradientAccumulation(gradient_accumulation)
-                inference_opts = deepcopy(opts)
-
-                # Initialize the dataset
-                num_batch = device_iterations * gradient_accumulation * num_replicate
-                data_size = num_batch * batch_size
-                dataset = self.TestDataset(
-                    labels=np.random.rand(data_size).astype(np.float32),
-                    node_features=[
-                        np.random.rand(node_feat_size).astype(np.float32) for ii in range(data_size)
-                    ],
-                    edge_features=[
-                        np.random.rand(edge_feat_size).astype(np.float32) for ii in range(data_size)
-                    ],
-                )
-
-                # Initialize the dataloader
-                train_dataloader = poptorch.DataLoader(
-                    options=training_opts,
-                    dataset=deepcopy(dataset),
-                    batch_size=batch_size,
-                    collate_fn=partial(global_batch_collator, batch_size),
-                )
-
-                val_dataloader = poptorch.DataLoader(
-                    options=inference_opts,
-                    dataset=deepcopy(dataset),
-                    batch_size=batch_size,
-                    collate_fn=partial(global_batch_collator, batch_size),
-                )
-
-                # Build the model, and run it on "IPU"
-                model = self.TestSimpleLightning(batch_size, node_feat_size, edge_feat_size, num_batch)
-
-                strategy = IPUStrategy(
-                    training_opts=training_opts, inference_opts=inference_opts, autoreport=True
-                )
-                trainer = Trainer(
-                    logger=True,
-                    enable_checkpointing=False,
-                    max_epochs=2,
-                    strategy=strategy,
-                    num_sanity_val_steps=0,
-                    accelerator="ipu",
-                    devices=1,
-                )
-                trainer.fit(model=model, train_dataloaders=train_dataloader, val_dataloaders=val_dataloader)
-
-    @pytest.mark.skip
-    def test_poptorch_graphium_deviceiterations_gradient_accumulation_full(self):
-        """
-        Test the device-iterations and gradient accumulation in a way
-        that is very similar to the Graphium code
-        to make sure that the dataloader and models handle them correcly.
-        """
-        with patch("poptorch.ipuHardwareIsAvailable", return_value=True):
-            with patch("lightning_graphcore.accelerator._IPU_AVAILABLE", new=True):
-                try:
-                    import poptorch
-                except Exception as e:
-                    warn(f"Skipping this test because poptorch is not available.\n{e}")
-                    return
-
-                from lightning_graphcore import IPUStrategy
-                import lightning_graphcore
-
-                # Current library imports
-                from graphium.config._loader import (
-                    load_datamodule,
-                    load_metrics,
-                    load_architecture,
-                    load_accelerator,
-                    load_predictor,
-                    load_trainer,
-                )
-                from graphium.utils.safe_run import SafeRun
-
-                # Simplified testing config - reflecting the toymix requirements
-                CONFIG_FILE = "tests/config_test_ipu_dataloader_multitask.yaml"
-                with open(CONFIG_FILE, "r") as f:
-                    cfg = yaml.safe_load(f)
-
-                cfg, accelerator = load_accelerator(cfg)
-
-                # Load the datamodule, and prepare the data
-                datamodule = load_datamodule(cfg, accelerator_type=accelerator)
-                datamodule.prepare_data()
-                metrics = load_metrics(cfg)
-                model_class, model_kwargs = load_architecture(cfg, in_dims=datamodule.in_dims)
-                # datamodule.setup()
-                predictor = load_predictor(
-                    cfg,
-                    model_class,
-                    model_kwargs,
-                    metrics,
-                    datamodule.get_task_levels(),
-                    accelerator,
-                    datamodule.featurization,
-                    datamodule.task_norms,
-                )
-                assert poptorch.ipuHardwareIsAvailable()
-                trainer = load_trainer(cfg, "test", accelerator, "date_time_suffix")
-                # Run the model training
-                with SafeRun(
-                    name="TRAINING", raise_error=cfg["constants"]["raise_train_error"], verbose=True
-                ):
-                    trainer.fit(model=predictor, datamodule=datamodule)
-
-
-if __name__ == "__main__":
-    ut.main()
diff --git a/tests/test_ipu_losses.py b/tests/test_ipu_losses.py
deleted file mode 100644
index cb18eee47..000000000
--- a/tests/test_ipu_losses.py
+++ /dev/null
@@ -1,172 +0,0 @@
-"""
---------------------------------------------------------------------------------
-Copyright (c) 2023 Valence Labs, Recursion Pharmaceuticals and Graphcore Limited.
-
-Use of this software is subject to the terms and conditions outlined in the LICENSE file.
-Unauthorized modification, distribution, or use is prohibited. Provided 'as is' without
-warranties of any kind.
-
-Valence Labs, Recursion Pharmaceuticals and Graphcore Limited are not liable for any damages arising from its use.
-Refer to the LICENSE file for the full terms and conditions.
---------------------------------------------------------------------------------
-"""
-
-
-import unittest as ut
-import torch
-from torch.nn import BCELoss, MSELoss, L1Loss, BCEWithLogitsLoss
-from copy import deepcopy
-import pytest
-
-from graphium.ipu.ipu_losses import BCELossIPU, MSELossIPU, L1LossIPU, BCEWithLogitsLossIPU, HybridCELossIPU
-from graphium.trainer.losses import HybridCELoss
-
-
-@pytest.mark.ipu
-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)
-
-    th = 0.7
-    nan_th = 0.2
-    preds_greater = preds > th
-    target_greater = (target > th).to(torch.float32)
-    target_greater_nan = deepcopy(target_greater)
-    is_nan = target < nan_th
-    target_greater_nan[target < nan_th] = torch.nan
-    target_nan = deepcopy(target)
-    target_nan[target < nan_th] = torch.nan
-
-    def test_bce(self):
-        preds = deepcopy(self.preds)
-        target = deepcopy(self.target_greater)
-        target_nan = deepcopy(self.target_greater_nan)
-
-        # Regular loss
-        loss_true = BCELoss()(preds, target)
-        loss_ipu = BCELossIPU()(preds, target)
-        self.assertFalse(loss_true.isnan(), "Regular BCELoss is NaN")
-        self.assertAlmostEqual(
-            loss_true.item(), loss_ipu.item(), places=6, msg="Regular BCELoss is different"
-        )
-
-        # Weighted loss
-        weight = torch.rand(preds.shape[1], dtype=torch.float32)
-        loss_true = BCELoss(weight=weight)(preds, target)
-        loss_ipu = BCELossIPU(weight=weight)(preds, target)
-        self.assertFalse(loss_true.isnan(), "Regular BCELoss is NaN")
-        self.assertAlmostEqual(loss_true.item(), loss_ipu.item(), msg="Weighted BCELoss is different")
-
-        # Regular loss with NaNs in target
-        not_nan = ~target_nan.isnan()
-        loss_true = BCELoss()(preds[not_nan], target[not_nan])
-        loss_ipu = BCELossIPU()(preds, target_nan)
-        self.assertFalse(loss_true.isnan(), "Regular BCELoss with target_nan is NaN")
-        self.assertFalse(loss_ipu.isnan(), "Regular BCELossIPU with target_nan is NaN")
-        self.assertAlmostEqual(
-            loss_true.item(), loss_ipu.item(), places=6, msg="Regular BCELoss with NaN is different"
-        )
-
-        # Weighted loss with NaNs in target
-        not_nan = ~target_nan.isnan()
-        weight = torch.rand(preds.shape, dtype=torch.float32)
-        loss_true = BCELoss(weight=weight[not_nan])(preds[not_nan], target_nan[not_nan])
-        loss_ipu = BCELossIPU(weight=weight)(preds, target_nan)
-        self.assertFalse(loss_true.isnan(), "Weighted BCELoss with target_nan is NaN")
-        self.assertFalse(loss_ipu.isnan(), "Weighted BCELossIPU with target_nan is NaN")
-        self.assertAlmostEqual(
-            loss_true.item(), loss_ipu.item(), places=6, msg="Weighted BCELoss with NaN is different"
-        )
-
-    def test_mse(self):
-        preds = deepcopy(self.preds)
-        target = deepcopy(self.target)
-        target_nan = deepcopy(self.target_nan)
-
-        # Regular loss
-        loss_true = MSELoss()(preds, target)
-        loss_ipu = MSELossIPU()(preds, target)
-        self.assertFalse(loss_true.isnan(), "Regular MSELoss is NaN")
-        self.assertAlmostEqual(
-            loss_true.item(), loss_ipu.item(), places=6, msg="Regular MSELoss is different"
-        )
-
-        # Regular loss with NaNs in target
-        not_nan = ~target_nan.isnan()
-        loss_true = MSELoss()(preds[not_nan], target[not_nan])
-        loss_ipu = MSELossIPU()(preds, target_nan)
-        self.assertFalse(loss_true.isnan(), "Regular MSELoss with target_nan is NaN")
-        self.assertFalse(loss_ipu.isnan(), "Regular MSELossIPU with target_nan is NaN")
-        self.assertAlmostEqual(
-            loss_true.item(), loss_ipu.item(), places=6, msg="Regular MSELoss with NaN is different"
-        )
-
-    def test_l1(self):
-        preds = deepcopy(self.preds)
-        target = deepcopy(self.target)
-        target_nan = deepcopy(self.target_nan)
-
-        # Regular loss
-        loss_true = L1Loss()(preds, target)
-        loss_ipu = L1LossIPU()(preds, target)
-        self.assertFalse(loss_true.isnan(), "Regular MAELoss is NaN")
-        self.assertAlmostEqual(
-            loss_true.item(), loss_ipu.item(), places=6, msg="Regular MAELoss is different"
-        )
-
-        # Regular loss with NaNs in target
-        not_nan = ~target_nan.isnan()
-        loss_true = L1Loss()(preds[not_nan], target[not_nan])
-        loss_ipu = L1LossIPU()(preds, target_nan)
-        self.assertFalse(loss_true.isnan(), "Regular MAELoss with target_nan is NaN")
-        self.assertFalse(loss_ipu.isnan(), "Regular MAELossIPU with target_nan is NaN")
-        self.assertAlmostEqual(
-            loss_true.item(), loss_ipu.item(), places=6, msg="Regular MAELoss with NaN is different"
-        )
-
-    def test_bce_logits(self):
-        preds = deepcopy(self.preds)
-        target = deepcopy(self.target_greater)
-        target_nan = deepcopy(self.target_greater_nan)
-
-        # Regular loss
-        loss_true = BCEWithLogitsLoss()(preds, target)
-        loss_ipu = BCEWithLogitsLossIPU()(preds, target)
-        self.assertFalse(loss_true.isnan(), "Regular BCEWithLogitsLoss is NaN")
-        self.assertAlmostEqual(
-            loss_true.item(), loss_ipu.item(), places=6, msg="Regular BCEWithLogitsLoss is different"
-        )
-
-        # Weighted loss
-        weight = torch.rand(preds.shape[1], dtype=torch.float32)
-        loss_true = BCEWithLogitsLoss(weight=weight)(preds, target)
-        loss_ipu = BCEWithLogitsLossIPU(weight=weight)(preds, target)
-        self.assertFalse(loss_true.isnan(), "Regular BCEWithLogitsLoss is NaN")
-        self.assertAlmostEqual(
-            loss_true.item(), loss_ipu.item(), msg="Weighted BCEWithLogitsLoss is different"
-        )
-
-        # Regular loss with NaNs in target
-        not_nan = ~target_nan.isnan()
-        loss_true = BCEWithLogitsLoss()(preds[not_nan], target[not_nan])
-        loss_ipu = BCEWithLogitsLossIPU()(preds, target_nan)
-        self.assertFalse(loss_true.isnan(), "Regular test_bce_logits with target_nan is NaN")
-        self.assertFalse(loss_ipu.isnan(), "Regular test_bce_logits with target_nan is NaN")
-        self.assertAlmostEqual(
-            loss_true.item(), loss_ipu.item(), places=6, msg="Regular BCELoss with NaN is different"
-        )
-
-        # Weighted loss with NaNs in target
-        not_nan = ~target_nan.isnan()
-        weight = torch.rand(preds.shape, dtype=torch.float32)
-        loss_true = BCEWithLogitsLoss(weight=weight[not_nan])(preds[not_nan], target_nan[not_nan])
-        loss_ipu = BCEWithLogitsLossIPU(weight=weight)(preds, target_nan)
-        self.assertFalse(loss_true.isnan(), "Weighted test_bce_logits with target_nan is NaN")
-        self.assertFalse(loss_ipu.isnan(), "Weighted test_bce_logits with target_nan is NaN")
-        self.assertAlmostEqual(
-            loss_true.item(),
-            loss_ipu.item(),
-            places=6,
-            msg="Weighted BCEWithLogitsLoss with NaN is different",
-        )
diff --git a/tests/test_ipu_options.py b/tests/test_ipu_options.py
deleted file mode 100644
index c3cc9aa3e..000000000
--- a/tests/test_ipu_options.py
+++ /dev/null
@@ -1,149 +0,0 @@
-"""
---------------------------------------------------------------------------------
-Copyright (c) 2023 Valence Labs, Recursion Pharmaceuticals and Graphcore Limited.
-
-Use of this software is subject to the terms and conditions outlined in the LICENSE file.
-Unauthorized modification, distribution, or use is prohibited. Provided 'as is' without
-warranties of any kind.
-
-Valence Labs, Recursion Pharmaceuticals and Graphcore Limited are not liable for any damages arising from its use.
-Refer to the LICENSE file for the full terms and conditions.
---------------------------------------------------------------------------------
-"""
-
-
-import pytest
-from graphium.config._loader import _get_ipu_opts, load_ipu_options
-from graphium.ipu.ipu_utils import ipu_options_list_to_file
-
-import tempfile
-from typing import Optional, List
-import os
-
-CONFIG_EXTRACT = {
-    "trainer": {"trainer": {"accumulate_grad_batches": 10}},
-    "accelerator": {
-        "type": "ipu",
-        "config_override": {
-            "datamodule": {
-                "args": {
-                    "ipu_dataloader_training_opts": {
-                        "mode": "async",
-                        "max_num_nodes_per_graph": 44,
-                        "max_num_edges_per_graph": 80,
-                    },
-                    "ipu_dataloader_inference_opts": {
-                        "mode": "async",
-                        "max_num_nodes_per_graph": 44,
-                        "max_num_edges_per_graph": 80,
-                    },
-                    "batch_size_training": 50,
-                    "batch_size_inference": 50,
-                }
-            },
-            "predictor": {"optim_kwargs": {"loss_scaling": 1024}},
-            "trainer": {"trainer": {"precision": 16, "accumulate_grad_batches": 4}},
-        },
-        "ipu_config": [
-            "deviceIterations(5)",
-            "replicationFactor(16)",
-            "TensorLocations.numIOTiles(128)",
-            '_Popart.set("defaultBufferingDepth", 128)',
-            "Precision.enableStochasticRounding(True)",
-        ],
-        "ipu_inference_config": [
-            "deviceIterations(1)",
-            "replicationFactor(4)",
-            "TensorLocations.numIOTiles(32)",
-            '_Popart.set("defaultBufferingDepth", 16)',
-            "Precision.enableStochasticRounding(True)",
-        ],
-    },
-}
-
-
-@pytest.mark.ipu
-def test_ipu_options():
-    try:
-        import poptorch
-
-        ipu_opts, ipu_inference_opts = _get_ipu_opts(CONFIG_EXTRACT)
-
-        # Define the expected IPU options for comparison
-        expected_ipu_opts = [
-            "deviceIterations(5)",
-            "replicationFactor(16)",
-            "TensorLocations.numIOTiles(128)",
-            '_Popart.set("defaultBufferingDepth", 128)',
-            "Precision.enableStochasticRounding(True)",
-        ]
-        expected_ipu_inference_opts = [
-            "deviceIterations(1)",
-            "replicationFactor(4)",
-            "TensorLocations.numIOTiles(32)",
-            '_Popart.set("defaultBufferingDepth", 16)',
-            "Precision.enableStochasticRounding(True)",
-        ]
-
-        # Test the _get_ipu_opts method
-        ipu_opts, ipu_inference_opts = _get_ipu_opts(CONFIG_EXTRACT)
-        assert ipu_opts == expected_ipu_opts, f"Expected {expected_ipu_opts}, but got {ipu_opts}"
-        assert (
-            ipu_inference_opts == expected_ipu_inference_opts
-        ), f"Expected {expected_ipu_inference_opts}, but got {ipu_inference_opts}"
-
-        # Test the load_ipu_options method
-        ipu_training_opts, ipu_inference_opts = load_ipu_options(
-            ipu_opts=ipu_opts,
-            seed=42,
-            model_name="test_model",
-            gradient_accumulation=CONFIG_EXTRACT["trainer"]["trainer"].get("accumulate_grad_batches", None),
-            ipu_inference_opts=ipu_inference_opts,
-        )
-
-        # Ensure that the options objects are not None
-        assert ipu_training_opts is not None, "Expected ipu_training_opts not to be None"
-        assert ipu_inference_opts is not None, "Expected ipu_inference_opts not to be None"
-
-        # Test the properties of the options objects
-        assert (
-            ipu_training_opts.replication_factor == 16
-        ), "Expected replication_factor of ipu_training_opts to be 16"
-        assert (
-            ipu_inference_opts.replication_factor == 4
-        ), "Expected replication_factor of ipu_inference_opts to be 4"
-        assert ipu_training_opts._popart, "Expected _popart of ipu_training_opts to be True"
-        assert ipu_inference_opts._popart, "Expected _popart of ipu_inference_opts to be True"
-
-    except ImportError:
-        pytest.skip("Skipping this test because poptorch is not available")
-
-
-@pytest.mark.ipu
-def test_ipu_options_list_to_file():
-    # Define a list of IPU options
-    ipu_options = [
-        "deviceIterations(5)",
-        "replicationFactor(16)",
-        "TensorLocations.numIOTiles(128)",
-        '_Popart.set("defaultBufferingDepth", 128)',
-        "Precision.enableStochasticRounding(True)",
-    ]
-
-    # Call the function with the list of IPU options
-    tmp_file = ipu_options_list_to_file(ipu_options)
-
-    # Check that the function returns a temporary file object
-    assert isinstance(tmp_file, tempfile._TemporaryFileWrapper)
-
-    # Check that the temporary file exists
-    assert os.path.exists(tmp_file.name)
-
-    # Check the contents of the temporary file
-    with open(tmp_file.name, "r") as f:
-        contents = f.read().splitlines()
-    assert contents == ipu_options
-
-    # Check the behavior when the input is None
-    tmp_file = ipu_options_list_to_file(None)
-    assert tmp_file is None
diff --git a/tests/test_ipu_poptorch.py b/tests/test_ipu_poptorch.py
deleted file mode 100644
index 4f951d504..000000000
--- a/tests/test_ipu_poptorch.py
+++ /dev/null
@@ -1,29 +0,0 @@
-"""
---------------------------------------------------------------------------------
-Copyright (c) 2023 Valence Labs, Recursion Pharmaceuticals and Graphcore Limited.
-
-Use of this software is subject to the terms and conditions outlined in the LICENSE file.
-Unauthorized modification, distribution, or use is prohibited. Provided 'as is' without
-warranties of any kind.
-
-Valence Labs, Recursion Pharmaceuticals and Graphcore Limited are not liable for any damages arising from its use.
-Refer to the LICENSE file for the full terms and conditions.
---------------------------------------------------------------------------------
-"""
-
-
-import pytest
-
-
-@pytest.mark.ipu
-def test_poptorch():
-    # Run this test only if poptorch is available
-    # Primarily to test the install and SDK is correctly activated
-    try:
-        import poptorch
-
-        opts = poptorch.Options()
-
-    except ImportError:
-        raise ImportError
-    assert True
diff --git a/tests/test_ipu_to_dense_batch.py b/tests/test_ipu_to_dense_batch.py
deleted file mode 100644
index 157b3b647..000000000
--- a/tests/test_ipu_to_dense_batch.py
+++ /dev/null
@@ -1,133 +0,0 @@
-"""
---------------------------------------------------------------------------------
-Copyright (c) 2023 Valence Labs, Recursion Pharmaceuticals and Graphcore Limited.
-
-Use of this software is subject to the terms and conditions outlined in the LICENSE file.
-Unauthorized modification, distribution, or use is prohibited. Provided 'as is' without
-warranties of any kind.
-
-Valence Labs, Recursion Pharmaceuticals and Graphcore Limited are not liable for any damages arising from its use.
-Refer to the LICENSE file for the full terms and conditions.
---------------------------------------------------------------------------------
-"""
-
-
-import pytest
-import torch
-from torch_geometric.data import Data, Batch
-from graphium.ipu.to_dense_batch import to_dense_batch
-
-import torch
-
-# Current library imports
-from graphium.config._loader import load_datamodule, load_metrics, load_architecture, load_accelerator
-
-
-@pytest.mark.ipu
-class test_IPUBatch:
-    @pytest.fixture(autouse=True)
-    def setup_class(self):
-        self.in_dim = 12
-        self.out_dim = 12
-        self.in_dim_edges = 10
-        self.out_dim_edges = 10
-        self.edge_idx1 = torch.stack(
-            [torch.tensor([0, 1, 2, 3, 2], dtype=torch.int), torch.tensor([1, 2, 3, 0, 0], dtype=torch.int)]
-        )
-        self.edge_idx2 = torch.stack(
-            [torch.tensor([0, 0, 0, 1], dtype=torch.int), torch.tensor([0, 1, 2, 0], dtype=torch.int)]
-        )
-        self.x1 = torch.randn(self.edge_idx1.max().item() + 1, self.in_dim, dtype=torch.float32)
-        self.e1 = torch.randn(self.edge_idx1.shape[-1], self.in_dim_edges, dtype=torch.float32)
-        self.x2 = torch.randn(self.edge_idx2.max().item() + 1, self.in_dim, dtype=torch.float32)
-        self.e2 = torch.randn(self.edge_idx2.shape[-1], self.in_dim_edges, dtype=torch.float32)
-        self.g1 = Data(feat=self.x1, edge_index=self.edge_idx1, edge_feat=self.e1)
-        self.g2 = Data(feat=self.x2, edge_index=self.edge_idx2, edge_feat=self.e2)
-        self.bg = Batch.from_data_list([self.g1, self.g2])
-        self.attn_kwargs = {"embed_dim": self.in_dim, "num_heads": 2, "batch_first": True}
-
-    # @pytest.mark.skip
-    @pytest.mark.parametrize("max_num_nodes_per_graph, batch_size", [(10, 5), (20, 10), (30, 15)])
-    def test_ipu_to_dense_batch(self, max_num_nodes_per_graph, batch_size):
-        # Run this test only if poptorch is available
-        try:
-            import poptorch
-
-            opts = poptorch.Options()
-            opts.useIpuModel(True)
-
-            class MyModel(torch.nn.Module):
-                def __init__(self):
-                    super(MyModel, self).__init__()
-
-                def forward(self, x, batch):
-                    return to_dense_batch(
-                        x,
-                        batch=batch,
-                        batch_size=batch_size,
-                        max_num_nodes_per_graph=max_num_nodes_per_graph,
-                        drop_nodes_last_graph=False,
-                    )
-
-            model = MyModel()
-            model = model.eval()
-            poptorch_model_inf = poptorch.inferenceModel(model, options=opts)
-            # for data in train_dataloader:
-            out, mask, idx = poptorch_model_inf(self.bg.feat, self.bg.batch)
-            # Check the output sizes
-            assert out.size() == torch.Size([batch_size, max_num_nodes_per_graph, 12])
-            # Check the mask for true / false values
-            assert mask.size() == torch.Size([batch_size, max_num_nodes_per_graph])
-            assert torch.sum(mask) == 7
-            assert (mask[0][:4] == True).all()
-            assert (mask[0][4:] == False).all()
-            assert (mask[1][:3] == True).all()
-            assert (mask[1][3:] == False).all()
-            assert (mask[2:] == False).all()
-
-            # Check the idx are all the true values in the mask
-            assert (mask.flatten()[idx] == True).all()
-            poptorch_model_inf.detachFromDevice()
-        except ImportError:
-            pytest.skip("Skipping this test because poptorch is not available")
-
-    def test_ipu_to_dense_batch_no_batch_no_max_nodes(self):
-        h_dense, mask = to_dense_batch(
-            self.bg.feat,
-            batch=None,
-            batch_size=None,
-            max_num_nodes_per_graph=None,
-            drop_nodes_last_graph=False,
-        )
-        # Add assertions to check the output as needed
-        assert torch.allclose(h_dense, self.bg.feat.unsqueeze(0), atol=1e-5), "Tensors are not equal"
-        assert mask.size(1) == h_dense.size(1)
-        assert mask.all().item(), "Not all values in the tensor are True"
-
-    def test_ipu_to_dense_batch_no_batch(self):
-        max_nodes_per_graph = 10
-        h_dense, mask, id = to_dense_batch(
-            self.bg.feat,
-            batch=None,
-            batch_size=None,
-            max_num_nodes_per_graph=max_nodes_per_graph,
-            drop_nodes_last_graph=False,
-        )
-        assert mask.size() == (1, max_nodes_per_graph)
-        assert torch.sum(mask) == 7
-        assert torch.equal(id, torch.arange(7))
-        assert h_dense.size() == (1, max_nodes_per_graph, self.bg.feat.size(-1))
-
-    def test_ipu_to_dense_batch_drop_last(self):
-        out, mask, idx = to_dense_batch(
-            self.bg.feat,
-            batch=None,
-            batch_size=None,
-            max_num_nodes_per_graph=3,
-            drop_nodes_last_graph=True,
-        )
-        # Add assertions to check the output as needed
-        assert mask.size(1) == out.size(1)
-        # Check the mask and output have been clipped
-        assert mask.size() == torch.Size([1, 3])
-        assert mask.all().item(), "Not all values in the tensor are True"
diff --git a/tests/test_mup.py b/tests/test_mup.py
index b60e0ccf3..f5fd48987 100644
--- a/tests/test_mup.py
+++ b/tests/test_mup.py
@@ -151,7 +151,7 @@ def test_feedforwardgraph_mup(self):
 
     def test_fullgraphmultitasknetwork(self):
         # Load the configuration file for the model
-        CONFIG_FILE = "tests/config_test_ipu_dataloader.yaml"
+        CONFIG_FILE = "tests/config_test_dataloader.yaml"
         with open(CONFIG_FILE, "r") as f:
             cfg = yaml.safe_load(f)
 
diff --git a/tests/test_packing.py b/tests/test_packing.py
deleted file mode 100644
index 3b378214b..000000000
--- a/tests/test_packing.py
+++ /dev/null
@@ -1,234 +0,0 @@
-"""
---------------------------------------------------------------------------------
-Copyright (c) 2023 Valence Labs, Recursion Pharmaceuticals and Graphcore Limited.
-
-Use of this software is subject to the terms and conditions outlined in the LICENSE file.
-Unauthorized modification, distribution, or use is prohibited. Provided 'as is' without
-warranties of any kind.
-
-Valence Labs, Recursion Pharmaceuticals and Graphcore Limited are not liable for any damages arising from its use.
-Refer to the LICENSE file for the full terms and conditions.
---------------------------------------------------------------------------------
-"""
-
-
-# General imports
-import unittest as ut
-import numpy as np
-
-import torch
-from torch_geometric.data import Data, Batch
-
-# Current library imports
-from graphium.utils.packing import (
-    smart_packing,
-    get_pack_sizes,
-    fast_packing,
-    hybrid_packing,
-    node_to_pack_indices_mask,
-)
-
-
-def random_packing(num_nodes, batch_size):
-    ipu_batch_size = int(len(num_nodes) / batch_size)
-    indices = np.arange(len(num_nodes))
-    np.random.shuffle(indices)
-    indices = np.reshape(indices, (ipu_batch_size, batch_size)).tolist()
-    return indices
-
-
-class test_Packing(ut.TestCase):
-    def test_smart_packing(self):
-        np.random.seed(42)
-
-        batch_sizes = [2, 4, 8, 16, 32, 64]
-        ipu_batch_sizes = [2, 3, 4, 8, 16, 32, 64]
-
-        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
-                global_batch = batch_size * ipu_batch_size
-                num_nodes = np.abs(np.random.gamma(2, 20, size=global_batch)).astype(int)
-
-                # Use the smart packing
-                packed_indices = smart_packing(num_nodes=num_nodes, batch_size=batch_size)
-                pack_num_nodes = get_pack_sizes(packed_indices, num_nodes)
-
-                # Use the random packing
-                rand_packed_indices = random_packing(num_nodes=num_nodes, batch_size=batch_size)
-                rand_pack_num_nodes = get_pack_sizes(rand_packed_indices, num_nodes)
-
-                # Assert that the smart packing is better than the random packing
-                self.assertLessEqual(max(pack_num_nodes), max(rand_pack_num_nodes), msg=err_msg)
-                self.assertGreaterEqual(min(pack_num_nodes), min(rand_pack_num_nodes), msg=err_msg)
-
-                # Assert that the total number of atoms is right
-                self.assertEqual(sum(pack_num_nodes), sum(num_nodes), msg=err_msg)
-                self.assertEqual(sum(rand_pack_num_nodes), sum(num_nodes), msg=err_msg)
-
-                # Assert that all index are there
-                self.assertListEqual(
-                    np.sort(np.asarray(packed_indices).flatten()).tolist(), np.arange(len(num_nodes)).tolist()
-                )
-                self.assertListEqual(
-                    np.sort(np.asarray(rand_packed_indices).flatten()).tolist(),
-                    np.arange(len(num_nodes)).tolist(),
-                )
-
-    def test_fast_packing(self):
-        np.random.seed(42)
-
-        # Start at 4 for fast_packing for better statistical significance
-        batch_sizes = [4, 8, 16, 32, 64]
-        ipu_batch_sizes = [4, 8, 16, 32, 64]
-
-        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
-                global_batch = batch_size * ipu_batch_size
-                num_nodes = np.abs(np.random.gamma(2, 20, size=global_batch)).astype(int)
-
-                # Use the smart packing
-                packed_indices = fast_packing(num_nodes=num_nodes, batch_size=batch_size)
-                pack_num_nodes = get_pack_sizes(packed_indices, num_nodes)
-
-                # Use the random packing
-                rand_packed_indices = random_packing(num_nodes=num_nodes, batch_size=batch_size)
-                rand_pack_num_nodes = get_pack_sizes(rand_packed_indices, num_nodes)
-
-                # Assert that the smart packing is better than the random packing
-                self.assertLessEqual(max(pack_num_nodes), max(rand_pack_num_nodes), msg=err_msg)
-                self.assertGreaterEqual(min(pack_num_nodes), min(rand_pack_num_nodes), msg=err_msg)
-
-                # Assert that the total number of atoms is right
-                self.assertEqual(sum(pack_num_nodes), sum(num_nodes), msg=err_msg)
-                self.assertEqual(sum(rand_pack_num_nodes), sum(num_nodes), msg=err_msg)
-
-                # Assert that all index are there
-                self.assertListEqual(
-                    np.sort(np.asarray(packed_indices).flatten()).tolist(), np.arange(len(num_nodes)).tolist()
-                )
-                self.assertListEqual(
-                    np.sort(np.asarray(rand_packed_indices).flatten()).tolist(),
-                    np.arange(len(num_nodes)).tolist(),
-                )
-
-    def test_hybrid_packing(self):
-        np.random.seed(42)
-
-        batch_sizes = [2, 4, 8, 16, 32, 64]
-        ipu_batch_sizes = [2, 3, 4, 8, 16, 32, 64]
-
-        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
-                global_batch = batch_size * ipu_batch_size
-                num_nodes = np.abs(np.random.gamma(2, 20, size=global_batch)).astype(int)
-
-                # Use the smart packing
-                packed_indices = hybrid_packing(num_nodes=num_nodes, batch_size=batch_size)
-                pack_num_nodes = get_pack_sizes(packed_indices, num_nodes)
-
-                # Use the random packing
-                rand_packed_indices = random_packing(num_nodes=num_nodes, batch_size=batch_size)
-                rand_pack_num_nodes = get_pack_sizes(rand_packed_indices, num_nodes)
-
-                # Assert that the smart packing is better than the random packing
-                self.assertLessEqual(max(pack_num_nodes), max(rand_pack_num_nodes), msg=err_msg)
-                self.assertGreaterEqual(min(pack_num_nodes), min(rand_pack_num_nodes), msg=err_msg)
-
-                # Assert that the total number of atoms is right
-                self.assertEqual(sum(pack_num_nodes), sum(num_nodes), msg=err_msg)
-                self.assertEqual(sum(rand_pack_num_nodes), sum(num_nodes), msg=err_msg)
-
-                # Assert that all index are there
-                self.assertListEqual(
-                    np.sort(np.asarray(packed_indices).flatten()).tolist(), np.arange(len(num_nodes)).tolist()
-                )
-                self.assertListEqual(
-                    np.sort(np.asarray(rand_packed_indices).flatten()).tolist(),
-                    np.arange(len(num_nodes)).tolist(),
-                )
-
-    def test_node_to_pack_indices_mask(self):
-        # Create a dummy batch
-        in_dim = 7
-        in_dim_edges = 11
-        max_num_nodes_per_graph = 20
-        batch_size_per_pack = 5
-
-        torch.manual_seed(42)
-
-        # Create a dummy batch of graphs
-        batch, all_num_nodes = [], []
-        for ii in range(100):
-            num_nodes = torch.randint(1, max_num_nodes_per_graph, (1,)).item()
-            all_num_nodes.append(num_nodes)
-            num_edges = abs(round(2.2 * num_nodes) + torch.randint(-2, 2, (1,)).item()) + 1
-            x = torch.randn(num_nodes, in_dim, dtype=torch.float32)
-            edge_idx = torch.randint(0, num_nodes, (2, num_edges))
-            e = torch.randn(edge_idx.shape[-1], in_dim_edges, dtype=torch.float32)
-            g = Data(h=x, edge_index=edge_idx, edge_attr=e)
-            batch.append(g)
-        batch = Batch.from_data_list(batch)
-
-        # Get the packing
-        packed_graph_idx = fast_packing(all_num_nodes, batch_size_per_pack)
-        pack_sizes = get_pack_sizes(packed_graph_idx, all_num_nodes)
-        max_pack_size = max(pack_sizes)
-        num_packs = len(pack_sizes)
-
-        # Get the node to pack indices and the mask
-        pack_from_node_idx, pack_attn_mask = node_to_pack_indices_mask(packed_graph_idx, all_num_nodes)
-
-        # Assert that the nodes to pack indices are correct
-        h = torch.arange(batch.num_nodes, dtype=torch.float32)
-        packed_shape = [num_packs, max_pack_size]
-        h_packed = torch.zeros(packed_shape)
-        h_packed[pack_from_node_idx[:, 0], pack_from_node_idx[:, 1]] = h
-        h_packed_unique = torch.sort(torch.unique(h_packed))[0]
-        np.testing.assert_array_equal(h_packed_unique, torch.arange(batch.num_nodes))
-        self.assertEqual(h_packed.sum(), h.sum())
-
-        # Test again with additional h dimension
-        h = batch.h
-        packed_shape = [num_packs, max_pack_size] + list(h.shape[1:])
-        h_packed = torch.zeros(packed_shape)
-        h_packed[pack_from_node_idx[:, 0], pack_from_node_idx[:, 1]] = h
-        h_packed_unique = torch.sort(torch.unique(h_packed))[0]
-        h_packed_unique = h_packed_unique[h_packed_unique != 0]
-        np.testing.assert_array_almost_equal(h_packed_unique, torch.unique(h))
-        self.assertAlmostEqual(h_packed.sum().item(), h.sum().item(), places=3)
-
-        # Assert that the mask is correct by counting the number of False values (the sum of squared number of nodes per pack)
-        num_false = (~pack_attn_mask).sum([1, 2])
-        num_expected = torch.as_tensor(
-            [sum([all_num_nodes[graph_idx] ** 2 for graph_idx in pack]) for pack in packed_graph_idx]
-        )
-        np.testing.assert_array_equal(num_false, num_expected)
-
-        # Assert that the mask is correct by counting the number of elements in each row and column
-        num_expected = []
-        for pack in packed_graph_idx:
-            pack_num_expected = []
-            for graph_idx in pack:
-                num_nodes = all_num_nodes[graph_idx]
-                for ii in range(num_nodes):
-                    pack_num_expected.append(num_nodes)
-            pack_num_expected.extend([0] * (max_pack_size - len(pack_num_expected)))
-            num_expected.append(pack_num_expected)
-        num_expected = torch.as_tensor(num_expected)
-        num_false_row = (~pack_attn_mask).sum([2])
-        num_false_col = (~pack_attn_mask).sum([1])
-        np.testing.assert_array_equal(num_false_row, num_expected)
-        np.testing.assert_array_equal(num_false_col, num_expected)
-
-
-if __name__ == "__main__":
-    ut.main()
diff --git a/tests/test_pyg_layers.py b/tests/test_pyg_layers.py
index 03498eb35..52caec668 100644
--- a/tests/test_pyg_layers.py
+++ b/tests/test_pyg_layers.py
@@ -231,7 +231,6 @@ def test_pnamessagepassinglayer(self):
         self.assertEqual(bg2.feat.shape[1], self.out_dim * layer.out_dim_factor)
         self.assertTrue((bg2.edge_feat == self.bg.edge_feat).all)
 
-    @pytest.mark.skip_ipu
     def test_dimenetlayer(self):
         from graphium.nn.encoders.bessel_pos_encoder import BesselSphericalPosEncoder
 
@@ -311,7 +310,7 @@ def test_preprocess3Dfeaturelayer(self):
         # bias: [batch, num_heads, nodes, nodes]
         # node_feature: [total_nodes, embed_dim]
         bias, node_feature = layer.forward(
-            bg, max_num_nodes_per_graph=4, on_ipu=False, positions_3d_key="positions_3d"
+            bg, max_num_nodes_per_graph=4, positions_3d_key="positions_3d"
         )
         self.assertEqual(bias.size(), torch.Size([2, num_heads, 4, 4]))
         self.assertFalse(np.isnan(bias.detach().numpy()).any())
diff --git a/tests/test_training.py b/tests/test_training.py
index 789bbca3a..ab8648aad 100644
--- a/tests/test_training.py
+++ b/tests/test_training.py
@@ -92,9 +92,7 @@ def call_cli_with_overrides(self, acc_type: str, acc_prec: str) -> None:
             "trainer.trainer.check_val_every_n_epoch=1",
             f"trainer.trainer.precision={acc_prec}",
         ]
-        if acc_type == "ipu":
-            overrides.append("accelerator.ipu_config=['useIpuModel(True)']")
-            overrides.append("accelerator.ipu_inference_config=['useIpuModel(True)']")
+
         # Backup the original sys.argv
         original_argv = sys.argv.copy()
 
@@ -110,20 +108,6 @@ def call_cli_with_overrides(self, acc_type: str, acc_prec: str) -> None:
     def test_cpu_cli_training(self):
         self.call_cli_with_overrides("cpu", "32")
 
-    @pytest.mark.ipu
-    @pytest.mark.skip
-    def test_ipu_cli_training(self):
-        with ut.patch("poptorch.ipuHardwareIsAvailable", return_value=True):
-            with ut.patch("lightning_graphcore.accelerator._IPU_AVAILABLE", new=True):
-                import poptorch
-
-                assert poptorch.ipuHardwareIsAvailable()
-                from lightning_graphcore.accelerator import _IPU_AVAILABLE
-
-                assert _IPU_AVAILABLE is True
-                self.call_cli_with_overrides("ipu", "16-true")
-
-
 
 def initialize_hydra(config_path, job_name="app"):
     if GlobalHydra.instance().is_initialized():