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

Issue:
Fixed OGB test. Added black
## Changelogs

- _enumerate the changes of that PR._

---

_Checklist:_

- [ ] _Was this PR discussed in an issue? It is recommended to first discuss a new feature into a GitHub issue before opening a PR._
- [ ] _Add tests to cover the fixed bug(s) or the new introduced feature(s) (if appropriate)._
- [ ] _Update the API documentation is a new function is added, or an existing one is deleted._
- [ ] _Write concise and explanatory changelogs above._
- [ ] _If possible, assign one of the following labels to the PR: `feature`, `fix` or `test` (or ask a maintainer to do it for you)._

---

_discussion related to that PR_



Discussion / hints:

Fix patch:
diff --git a/expts/configs/config_tdc_admet_demo.yaml b/expts/configs/config_tdc_admet_demo.yaml
new file mode 100644
index 000000000..aac4d0e50
--- /dev/null
+++ b/expts/configs/config_tdc_admet_demo.yaml
@@ -0,0 +1,315 @@
+# Testing the gcn model with the PCQMv2 dataset on IPU.
+constants:
+  name: &name tdc_admet_demo
+  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
+
+datamodule:
+  module_type: "ADMETBenchmarkDataModule"
+  args:
+    # TDC specific
+    tdc_benchmark_names: null
+    tdc_train_val_seed: *seed
+    # Featurization
+    prepare_dict_or_graph: pyg:graph
+    featurization_n_jobs: 30
+    featurization_progress: True
+    featurization_backend: "loky"
+    processed_graph_data_path: "../datacache/tdc-admet-demo/"
+    featurization:
+      atom_property_list_onehot: [atomic-number, group, period, total-valence]
+      atom_property_list_float: [degree, formal-charge, radical-electron, aromatic, 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:
+        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.
+
+
+architecture:
+  model_type: FullGraphMultiTaskNetwork
+  mup_base_path: null
+  pre_nn:   # Set as null to avoid a pre-nn network
+    out_dim: 64
+    hidden_dims: 256
+    depth: 2
+    activation: relu
+    last_activation: none
+    dropout: &dropout 0.18
+    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: 64 # or otherwise the correct value
+    out_dim: &gnn_dim 96
+    hidden_dims: *gnn_dim
+    depth: 4
+    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:
+    caco2_wang: &regression_head
+      task_level: graph
+      out_dim: 1
+      hidden_dims: 64
+      depth: 2
+      activation: relu
+      last_activation: none
+      dropout: *dropout
+      normalization: *normalization
+      last_normalization: "none"
+      residual_type: none
+    hia_hou: &classification_head
+      task_level: graph
+      out_dim: 1
+      hidden_dims: 64
+      depth: 2
+      activation: relu
+      last_activation: sigmoid
+      dropout: *dropout
+      normalization: *normalization
+      last_normalization: "none"
+      residual_type: none
+    pgp_broccatelli: *classification_head
+    bioavailability_ma: *classification_head
+    lipophilicity_astrazeneca: *regression_head
+    solubility_aqsoldb: *regression_head
+    bbb_martins: *classification_head
+    ppbr_az: *regression_head
+    vdss_lombardo: *regression_head
+    cyp2d6_veith: *classification_head
+    cyp3a4_veith: *classification_head
+    cyp2c9_veith: *classification_head
+    cyp2d6_substrate_carbonmangels: *classification_head
+    cyp3a4_substrate_carbonmangels: *classification_head
+    cyp2c9_substrate_carbonmangels: *classification_head
+    half_life_obach: *regression_head
+    clearance_microsome_az: *regression_head
+    clearance_hepatocyte_az: *regression_head
+    herg: *classification_head
+    ames: *classification_head
+    dili: *classification_head
+    ld50_zhu: *regression_head
+
+#Task-specific
+predictor:
+  metrics_on_progress_bar:
+    # All below metrics are directly copied from the TDC website.
+    # For more information, see https://tdcommons.ai/benchmark/admet_group/overview/
+    caco2_wang: ["mae"]
+    hia_hou: ["auroc"]
+    pgp_broccatelli: ["auroc"]
+    bioavailability_ma: ["auroc"]
+    lipophilicity_astrazeneca: ["mae"]
+    solubility_aqsoldb: ["mae"]
+    bbb_martins: ["auroc"]
+    ppbr_az: ["mae"]
+    vdss_lombardo: ["spearman"]
+    cyp2d6_veith: ["auprc"]
+    cyp3a4_veith: ["auprc"]
+    cyp2c9_veith: ["auprc"]
+    cyp2d6_substrate_carbonmangels: ["auprc"]
+    cyp3a4_substrate_carbonmangels: ["auprc"]
+    cyp2c9_substrate_carbonmangels: ["auprc"]
+    half_life_obach: ["spearman"]
+    clearance_microsome_az: ["spearman"]
+    clearance_hepatocyte_az: ["spearman"]
+    herg: ["mae"]
+    ames: ["auroc"]
+    dili: ["auroc"]
+    ld50_zhu: ["auroc"]
+  loss_fun:
+    caco2_wang: mae
+    hia_hou: bce
+    pgp_broccatelli: bce
+    bioavailability_ma: bce
+    lipophilicity_astrazeneca: mae
+    solubility_aqsoldb: mae
+    bbb_martins: bce
+    ppbr_az: mae
+    vdss_lombardo: mae
+    cyp2d6_veith: bce
+    cyp3a4_veith: bce
+    cyp2c9_veith: bce
+    cyp2d6_substrate_carbonmangels: bce
+    cyp3a4_substrate_carbonmangels: bce
+    cyp2c9_substrate_carbonmangels: bce
+    half_life_obach: mae
+    clearance_microsome_az: mae
+    clearance_hepatocyte_az: mae
+    herg: bce
+    ames: bce
+    dili: bce
+    ld50_zhu: mae
+  random_seed: *seed
+  optim_kwargs:
+    lr: 4.e-5 # warmup can be scheduled using torch_scheduler_kwargs
+  torch_scheduler_kwargs:
+    module_type: WarmUpLinearLR
+    max_num_epochs: &max_epochs 10
+    warmup_epochs: 10
+    verbose: False
+  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:
+  caco2_wang: &regression_metrics
+    - name: mae
+      metric: mae
+      target_nan_mask: null
+      multitask_handling: flatten
+      threshold_kwargs: null
+    - name: spearman
+      metric: spearmanr
+      threshold_kwargs: null
+      target_nan_mask: null
+      multitask_handling: mean-per-label
+    - name: pearson
+      metric: pearsonr
+      threshold_kwargs: null
+      target_nan_mask: null
+      multitask_handling: mean-per-label
+    - name: r2_score
+      metric: r2
+      target_nan_mask: null
+      multitask_handling: mean-per-label
+      threshold_kwargs: null
+  hia_hou: &classification_metrics
+    - name: auroc
+      metric: auroc
+      task: binary
+      multitask_handling: mean-per-label
+      threshold_kwargs: null
+    - name: auprc
+      metric: average_precision
+      task: binary
+      multitask_handling: mean-per-label
+      threshold_kwargs: null
+    - name: accuracy
+      metric: accuracy
+      multitask_handling: mean-per-label
+      target_to_int: True
+      average: micro
+      threshold_kwargs: &threshold_05
+        operator: greater
+        threshold: 0.5
+        th_on_preds: True
+        th_on_target: True
+    - name: mcc
+      metric: mcc
+      num_classes: 2
+      multitask_handling: mean-per-label
+      target_to_int: True
+      average: micro
+      threshold_kwargs: *threshold_05
+  pgp_broccatelli: *classification_metrics
+  bioavailability_ma: *classification_metrics
+  lipophilicity_astrazeneca: *regression_metrics
+  solubility_aqsoldb: *regression_metrics
+  bbb_martins: *classification_metrics
+  ppbr_az: *regression_metrics
+  vdss_lombardo: *regression_metrics
+  cyp2d6_veith: *classification_metrics
+  cyp3a4_veith: *classification_metrics
+  cyp2c9_veith: *classification_metrics
+  cyp2d6_substrate_carbonmangels: *classification_metrics
+  cyp3a4_substrate_carbonmangels: *classification_metrics
+  cyp2c9_substrate_carbonmangels: *classification_metrics
+  half_life_obach: *regression_metrics
+  clearance_microsome_az: *regression_metrics
+  clearance_hepatocyte_az: *regression_metrics
+  herg: *classification_metrics
+  ames: *classification_metrics
+  dili: *classification_metrics
+  ld50_zhu: *regression_metrics
+trainer:
+  seed: *seed
+  logger:
+    save_dir: logs/tdc-admet-demo/
+    name: *name
+    project: *name
+  model_checkpoint:
+    dirpath: models_checkpoints/tdc-admet-demo/
+    filename: *name
+    save_last: True
+  trainer:
+    max_epochs: *max_epochs
+    min_epochs: 1
+    check_val_every_n_epoch: 20
diff --git a/graphium/data/__init__.py b/graphium/data/__init__.py
index af1dbe131..0a8fcd24d 100644
--- a/graphium/data/__init__.py
+++ b/graphium/data/__init__.py
@@ -5,6 +5,7 @@
 
 from .datamodule import GraphOGBDataModule
 from .datamodule import MultitaskFromSmilesDataModule
+from .datamodule import ADMETBenchmarkDataModule
 from .datamodule import FakeDataModule
 
 from .dataset import SingleTaskDataset
diff --git a/graphium/data/datamodule.py b/graphium/data/datamodule.py
index d598fbb40..2ec328fad 100644
--- a/graphium/data/datamodule.py
+++ b/graphium/data/datamodule.py
@@ -1,14 +1,16 @@
+import tempfile
+from contextlib import redirect_stderr, redirect_stdout
 from typing import Type, List, Dict, Union, Any, Callable, Optional, Tuple, Iterable, Literal
 
 import os
-import glob
 from functools import partial
 import importlib.resources
 import zipfile
 from copy import deepcopy
 import time
 import gc
-from rdkit import Chem
+
+import platformdirs
 import re
 from graphium.data.utils import get_keys
 
@@ -29,7 +31,6 @@
 from lightning.pytorch.trainer.states import RunningStage
 
 import torch
-from torch_geometric.data import Data
 from torch.utils.data.dataloader import DataLoader, Dataset
 from torch.utils.data import Subset
 
@@ -642,14 +643,15 @@ def get_max_num_edges_datamodule(self, stages: Optional[List[str]] = None) -> in
 class DatasetProcessingParams:
     def __init__(
         self,
-        task_level: str = None,
-        df: pd.DataFrame = None,
+        task_level: Optional[str] = None,
+        df: Optional[pd.DataFrame] = None,
         df_path: Optional[Union[str, os.PathLike]] = None,
-        smiles_col: str = None,
+        smiles_col: Optional[str] = None,
         label_cols: List[str] = None,
-        weights_col: str = None,  # Not needed
-        weights_type: str = None,  # Not needed
-        idx_col: str = None,
+        weights_col: Optional[str] = None,  # Not needed
+        weights_type: Optional[str] = None,  # Not needed
+        idx_col: Optional[str] = None,
+        mol_ids_col: Optional[str] = None,
         sample_size: Union[int, float, Type[None]] = None,
         split_val: float = 0.2,
         split_test: float = 0.2,
@@ -669,12 +671,17 @@ def __init__(
             weights_col: The column name of the weights
             weights_type: The type of weights
             idx_col: The column name of the indices
+            mol_ids_col: The column name of the molecule ids
             sample_size: The size of the sample
             split_val: The fraction of the data to use for validation
             split_test: The fraction of the data to use for testing
             seed: The seed to use for the splits and subsampling
             splits_path: The path to the splits
         """
+
+        if df is None and df_path is None:
+            raise ValueError("Either `df` or `df_path` must be provided")
+
         self.df = df
         self.task_level = task_level
         self.df_path = df_path
@@ -683,6 +690,7 @@ def __init__(
         self.weights_col = weights_col
         self.weights_type = weights_type
         self.idx_col = idx_col
+        self.mol_ids_col = mol_ids_col
         self.sample_size = sample_size
         self.split_val = split_val
         self.split_test = split_test
@@ -752,7 +760,7 @@ def _dataloader(self, dataset: Dataset, **kwargs) -> "poptorch.DataLoader":
 class MultitaskFromSmilesDataModule(BaseDataModule, IPUDataModuleModifier):
     def __init__(
         self,
-        task_specific_args: Dict[str, Any],  # TODO: Replace this with DatasetParams
+        task_specific_args: Union[DatasetProcessingParams, Dict[str, Any]],
         cache_data_path: Optional[Union[str, os.PathLike]] = None,
         processed_graph_data_path: Optional[Union[str, os.PathLike]] = None,
         featurization: Optional[Union[Dict[str, Any], omegaconf.DictConfig]] = None,
@@ -859,11 +867,16 @@ def __init__(
 
         self.task_specific_args = task_specific_args
 
-        # TODO: Have the input argument to the Data Module be of type DatasetParams
-        self.task_dataset_processing_params = {
-            self._get_task_key(ds_args["task_level"], task): DatasetProcessingParams(**ds_args)
-            for task, ds_args in task_specific_args.items()
-        }
+        self.task_dataset_processing_params = {}
+        for task, ds_args in task_specific_args.items():
+            if not isinstance(ds_args, DatasetProcessingParams):
+                # This is needed as long as not all classes have been migrated
+                # to use the new `DatasetProcessingParams` class
+                ds_args = DatasetProcessingParams(**ds_args)
+
+            key = self._get_task_key(ds_args.task_level, task)
+            self.task_dataset_processing_params[key] = ds_args
+
         self.featurization_n_jobs = featurization_n_jobs
         self.featurization_progress = featurization_progress
         self.featurization_backend = featurization_backend
@@ -1034,15 +1047,17 @@ def has_atoms_after_h_removal(smiles):
             for count in range(len(dataset_args["smiles"])):
                 all_tasks.append(task)
         # Get all unique mol ids
-        all_mol_ids = smiles_to_unique_mol_ids(
+        all_unique_mol_ids = smiles_to_unique_mol_ids(
             all_smiles,
             n_jobs=self.featurization_n_jobs,
             featurization_batch_size=self.featurization_batch_size,
             backend=self.featurization_backend,
         )
-        unique_mol_ids, unique_idx, inv = np.unique(all_mol_ids, return_index=True, return_inverse=True)
+        _, unique_ids_idx, unique_ids_inv = np.unique(
+            all_unique_mol_ids, return_index=True, return_inverse=True
+        )
 
-        smiles_to_featurize = [all_smiles[ii] for ii in unique_idx]
+        smiles_to_featurize = [all_smiles[ii] for ii in unique_ids_idx]
 
         # Convert SMILES to features
         features, _ = self._featurize_molecules(smiles_to_featurize)
@@ -1052,7 +1067,7 @@ def has_atoms_after_h_removal(smiles):
             task_dataset_args[task]["features"] = []
             task_dataset_args[task]["idx_none"] = []
         # Create a list of features matching up with the original smiles
-        all_features = [features[unique_idx] for unique_idx in inv]
+        all_features = [features[unique_idx] for unique_idx in unique_ids_inv]
 
         # Add the features to the task-specific data
         for all_idx, task in enumerate(all_tasks):
@@ -1066,7 +1081,7 @@ def has_atoms_after_h_removal(smiles):
             for idx, feat in enumerate(args["features"]):
                 if did_featurization_fail(feat):
                     idx_none.append(idx)
-            this_unique_ids = all_mol_ids[idx_per_task[task][0] : idx_per_task[task][1]]
+            this_unique_ids = all_unique_mol_ids[idx_per_task[task][0] : idx_per_task[task][1]]
             df, features, smiles, labels, sample_idx, extras, this_unique_ids = self._filter_none_molecules(
                 idx_none,
                 task_df[task],
@@ -1089,6 +1104,7 @@ def has_atoms_after_h_removal(smiles):
                 labels=task_dataset_args[task]["labels"],
                 smiles=task_dataset_args[task]["smiles"],
                 unique_ids=this_unique_ids,
+                indices=task_dataset_args[task]["sample_idx"],
                 **task_dataset_args[task]["extras"],
             )
 
@@ -1712,11 +1728,12 @@ def _extract_smiles_labels(
         self,
         df: pd.DataFrame,
         task_level: str,
-        smiles_col: str = None,
+        smiles_col: Optional[str] = None,
         label_cols: List[str] = [],
-        idx_col: str = None,
-        weights_col: str = None,
-        weights_type: str = None,
+        idx_col: Optional[str] = None,
+        mol_ids_col: Optional[str] = None,
+        weights_col: Optional[str] = None,
+        weights_type: Optional[str] = None,
     ) -> Tuple[
         np.ndarray, np.ndarray, Union[Type[None], np.ndarray], Dict[str, Union[Type[None], np.ndarray]]
     ]:
@@ -1729,6 +1746,7 @@ def _extract_smiles_labels(
             smiles_col: Name of the column containing the SMILES
             label_cols: List of column names containing the labels
             idx_col: Name of the column containing the index
+            mol_ids_col: Name of the column containing the molecule ids
             weights_col: Name of the column containing the weights
             weights_type: Type of weights to use.
         Returns:
@@ -1765,12 +1783,16 @@ def _extract_smiles_labels(
         else:
             labels = float("nan") + np.zeros([len(smiles), 0])
 
-        indices = None  # What are indices for?
+        # Get the indices, used for sub-sampling and splitting the dataset
         if idx_col is not None:
-            indices = df[idx_col].values
-
+            df = df.set_index(idx_col)
         sample_idx = df.index.values
 
+        # Get the molecule ids
+        mol_ids = None
+        if mol_ids_col is not None:
+            mol_ids = df[mol_ids_col].values
+
         # Extract the weights
         weights = None
         if weights_col is not None:
@@ -1797,7 +1819,7 @@ def _extract_smiles_labels(
 
             weights /= np.max(weights)  # Put the max weight to 1
 
-        extras = {"indices": indices, "weights": weights}
+        extras = {"weights": weights, "mol_ids": mol_ids}
         return smiles, labels, sample_idx, extras
 
     def _get_split_indices(
@@ -1859,15 +1881,14 @@ def _get_split_indices(
             elif file_type in ["csv", "tsv"]:
                 with fsspec.open(str(splits_path)) as f:
                     splits = self._read_csv(splits_path)
-                splits = splits.dropna()
             else:
                 raise ValueError(
                     f"file type `{file_type}` for `{splits_path}` not recognised, please use .pt, .csv or .tsv"
                 )
             train, val, test = split_names
-            train_indices = np.asarray(splits[train]).astype("int").tolist()
-            val_indices = np.asarray(splits[val]).astype("int").tolist()
-            test_indices = np.asarray(splits[test]).astype("int").tolist()
+            train_indices = np.asarray(splits[train].dropna()).astype("int").tolist()
+            val_indices = np.asarray(splits[val].dropna()).astype("int").tolist()
+            test_indices = np.asarray(splits[test].dropna()).astype("int").tolist()
 
         # Filter train, val and test indices
         _, train_idx, _ = np.intersect1d(sample_idx, train_indices, return_indices=True)
@@ -2114,7 +2135,7 @@ def __repr__(self) -> str:
 class GraphOGBDataModule(MultitaskFromSmilesDataModule):
     def __init__(
         self,
-        task_specific_args: Dict[str, Dict[str, Any]],  # TODO: Replace this with DatasetParams
+        task_specific_args: Dict[str, Union[DatasetProcessingParams, Dict[str, Any]]],
         cache_data_path: Optional[Union[str, os.PathLike]] = None,
         featurization: Optional[Union[Dict[str, Any], omegaconf.DictConfig]] = None,
         batch_size_training: int = 16,
@@ -2172,12 +2193,12 @@ def __init__(
             # Get OGB metadata
             this_metadata = self._get_dataset_metadata(task_args["dataset_name"])
             # Get dataset
-            df, idx_col, smiles_col, label_cols, splits_path = self._load_dataset(
+            df, mol_ids_col, smiles_col, label_cols, splits_path = self._load_dataset(
                 this_metadata, sample_size=task_args.get("sample_size", None)
             )
             new_task_specific_args[task_name] = {
                 "df": df,
-                "idx_col": idx_col,
+                "mol_ids_col": mol_ids_col,
                 "smiles_col": smiles_col,
                 "label_cols": label_cols,
                 "splits_path": splits_path,
@@ -2231,7 +2252,7 @@ def _load_dataset(
                 meaning that all molecules will be considered.
         Returns:
             df: Pandas dataframe containing the dataset.
-            idx_col: Name of the column containing the molecule index.
+            mol_ids_col: Name of the column containing the molecule ids.
             smiles_col: Name of the column containing the SMILES.
             label_cols: List of column names containing the labels.
             splits_path: Path to the file containing the train/val/test splits.
@@ -2300,15 +2321,15 @@ def _load_dataset(
 
         # Get column names: OGB columns are predictable
         if metadata["download_name"].startswith("pcqm4m"):
-            idx_col = df.columns[0]
+            mol_ids_col = df.columns[0]
             smiles_col = df.columns[-2]
             label_cols = df.columns[-1:].to_list()
         else:
-            idx_col = df.columns[-1]
+            mol_ids_col = df.columns[-1]
             smiles_col = df.columns[-2]
             label_cols = df.columns[:-2].to_list()
 
-        return df, idx_col, smiles_col, label_cols, splits_path
+        return df, mol_ids_col, smiles_col, label_cols, splits_path
 
     def _get_dataset_metadata(self, dataset_name: str) -> Dict[str, Any]:
         ogb_metadata = self._get_ogb_metadata()
@@ -2337,6 +2358,171 @@ def _get_ogb_metadata(self):
         return ogb_metadata
 
 
+class ADMETBenchmarkDataModule(MultitaskFromSmilesDataModule):
+    """
+    Wrapper to use the ADMET benchmark group from the TDC (Therapeutics Data Commons).
+
+    !!! warning "Dependency"
+
+        This class requires [PyTDC](https://pypi.org/project/PyTDC/) to be installed.
+
+    !!! note "Citation"
+
+        Huang, K., Fu, T., Gao, W., Zhao, Y., Roohani, Y., Leskovec, J., Coley, C., Xiao, C., Sun, J., & Zitnik, M. (2021).
+        Therapeutics Data Commons: Machine Learning Datasets and Tasks for Drug Discovery and Development.
+        Proceedings of Neural Information Processing Systems, NeurIPS Datasets and Benchmarks.
+
+
+    Parameters:
+        tdc_benchmark_names: This can be any subset of the benchmark names that make up the ADMET benchmarking group.
+            If `None`, uses the complete benchmarking group. For all full list of options, see
+            [the TDC website](https://tdcommons.ai/benchmark/admet_group/overview/) or use:
+
+           ```python
+           import tdc.utils.retrieve_benchmark_names
+           retrieve_benchmark_names("admet_group")
+           ```
+        tdc_train_val_seed: TDC recommends a default splitting method for the train-val split. This parameter
+          is used to seed that splitting method.
+    """
+
+    def __init__(
+        self,
+        # TDC-specific
+        tdc_benchmark_names: Optional[Union[str, List[str]]] = None,
+        tdc_train_val_seed: int = 0,
+        # Inherited arguments from superclass
+        cache_data_path: Optional[Union[str, os.PathLike]] = None,
+        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,
+        featurization_n_jobs: int = -1,
+        featurization_progress: bool = False,
+        featurization_backend: str = "loky",
+        collate_fn: Optional[Callable] = None,
+        prepare_dict_or_graph: str = "pyg:graph",
+        **kwargs,
+    ):
+        try:
+            from tdc.benchmark_group import admet_group
+            from tdc.utils import retrieve_benchmark_names
+        except ImportError as error:
+            # To make sure we use the exact same train-test set and preprocessing as other benchmark entries,
+            # we rely on the PyTDC library.
+            raise RuntimeError(
+                f"To use {self.__class__.__name__}, `PyTDC` needs to be installed. "
+                f"Please install it with `pip install PyTDC`"
+            ) from error
+
+        # Pick a path to save the TDC data to
+        tdc_cache_dir = fs.get_cache_dir("tdc")
+        tdc_cache_dir = fs.join(tdc_cache_dir, "ADMET_Benchmark")
+        fs.mkdir(tdc_cache_dir, exist_ok=True)
+
+        # Create the benchmark group object
+        # NOTE (cwognum): We redirect stderr and stdout to a file since TDC uses print statements,
+        #  which quickly pollute the logs. Ideally, we would use `redirect_stderr(None)`, but that breaks TQDM.
+        with tempfile.TemporaryFile("w") as f:
+            with redirect_stderr(f):
+                with redirect_stdout(f):
+                    self.group = admet_group(path=tdc_cache_dir)
+
+        # By default, use all available benchmarks in a benchmark group
+        if tdc_benchmark_names is None:
+            tdc_benchmark_names = retrieve_benchmark_names("admet_group")
+        if isinstance(tdc_benchmark_names, str):
+            tdc_benchmark_names = [tdc_benchmark_names]
+
+        # Create the task-specific arguments
+        logger.info(
+            f"Preparing the TDC ADMET Benchmark Group splits for each of the {len(tdc_benchmark_names)} benchmarks."
+        )
+
+        task_specific_args = {
+            t: self._get_task_specific_arguments(t, tdc_train_val_seed, tdc_cache_dir)
+            for t in tdc_benchmark_names
+        }
+
+        super().__init__(
+            task_specific_args=task_specific_args,
+            cache_data_path=cache_data_path,
+            featurization=featurization,
+            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,
+            featurization_n_jobs=featurization_n_jobs,
+            featurization_progress=featurization_progress,
+            featurization_backend=featurization_backend,
+            collate_fn=collate_fn,
+            prepare_dict_or_graph=prepare_dict_or_graph,
+            **kwargs,
+        )
+
+    def _get_task_specific_arguments(self, name: str, seed: int, cache_dir: str) -> DatasetProcessingParams:
+        """
+        Loads the data and split from the TDC benchmark group object.
+
+        For the train-test split, this is fixed.
+
+        For the train-val split, this does not have to be fixed. Here we use the default splitting method
+        from TDC to do the train-val split and allow the seed to be changed. This is likely to best match
+        other entries in the benchmarking group.
+        """
+
+        benchmark = self.group.get(name)
+
+        # Get the default train-val-test split
+
+        # NOTE (cwognum): TDC prints by default to stderr, which pollutes the logs quite a bit.
+        #  This context manager mutes these by temporarily writing to a file.
+        #  Ideally, we would use `redirect_stderr(None)`, but that breaks TQDM.
+
+        with tempfile.TemporaryFile("w") as f:
+            with redirect_stderr(f):
+                train, val = self.group.get_train_valid_split(seed, name)
+        test = benchmark["test"]
+
+        # Convert to the Graphium format
+        n_val = len(val)
+        n_test = len(test)
+        n_train = len(train)
+        max_len = max(n_train, n_val, n_test)
+        total_len = n_train + n_val + n_test
+
+        data = pd.concat([train, val, test], ignore_index=True)
+
+        # NOTE (cwognum): We need to convert the labels to float, since we use NaNs down the line.
+        #  If you uncomment this line, collating the labels will raise an overflow by converting a NaN to the int dtype.
+        data["Y"] = data["Y"].astype(float)
+
+        split = pd.DataFrame(
+            {
+                "train": list(range(n_train)),
+                "val": list(range(n_train, n_train + n_val)) + [float("nan")] * (max_len - n_val),
+                "test": list(range(n_train + n_val, total_len)) + [float("nan")] * (max_len - n_test),
+            }
+        )
+        split_path = fs.join(cache_dir, f"{name}_split.csv")
+        split.to_csv(split_path, index=False)
+
+        return DatasetProcessingParams(
+            df=data,
+            idx_col=None,
+            smiles_col="Drug",
+            label_cols=["Y"],
+            splits_path=split_path,
+            split_names=["train", "val", "test"],
+            task_level="graph",
+        )
+
+
 class FakeDataModule(MultitaskFromSmilesDataModule):
     """
     A fake datamodule that generates artificial data by mimicking the true data coming
@@ -2471,7 +2657,7 @@ def prepare_data(self):
             idx_per_task[task] = (total_len, total_len + num_smiles)
             total_len += num_smiles
         # Get all unique mol ids
-        all_mol_ids = smiles_to_unique_mol_ids(
+        all_unique_mol_ids = smiles_to_unique_mol_ids(
             all_smiles,
             n_jobs=self.featurization_n_jobs,
             featurization_batch_size=self.featurization_batch_size,
@@ -2487,7 +2673,8 @@ def prepare_data(self):
                 features=task_dataset_args[task]["features"],
                 labels=task_dataset_args[task]["labels"],
                 smiles=task_dataset_args[task]["smiles"],
-                unique_ids=all_mol_ids[idx_per_task[task][0] : idx_per_task[task][1]],
+                indices=task_dataset_args[task]["sample_idx"],
+                unique_ids=all_unique_mol_ids[idx_per_task[task][0] : idx_per_task[task][1]],
                 **task_dataset_args[task]["extras"],
             )
 
diff --git a/graphium/data/dataset.py b/graphium/data/dataset.py
index e07e083ce..fc99b720b 100644
--- a/graphium/data/dataset.py
+++ b/graphium/data/dataset.py
@@ -22,9 +22,10 @@ def __init__(
         labels: List[Union[torch.Tensor, np.ndarray]],
         features: Optional[List[Union[Data, "GraphDict"]]] = None,
         smiles: Optional[List[str]] = None,
-        indices: Optional[List[str]] = None,
+        indices: Optional[List[int]] = None,
         weights: Optional[Union[torch.Tensor, np.ndarray]] = None,
         unique_ids: Optional[List[str]] = None,
+        mol_ids: Optional[List[str]] = None,
     ):
         r"""
         dataset for a single task
@@ -34,31 +35,25 @@ def __init__(
             smiles: A list of smiles
             indices: A list of indices
             weights: A list of weights
-            unique_ids: A list of unique ids
+            unique_ids: A list of unique ids for each molecule generated from `datamol.unique_id`
+            mol_ids: A list of ids coming from the original dataset. Useful to identify the molecule in the original dataset.
         """
 
         # Verify that all lists are the same length
         numel = len(labels)
-        if features is not None:
-            assert (
-                len(features) == numel
-            ), f"features must be the same length as labels, got {len(features)} and {numel}"
-        if smiles is not None:
-            assert (
-                len(smiles) == numel
-            ), f"smiles must be the same length as labels, got {len(smiles)} and {numel}"
-        if indices is not None:
-            assert (
-                len(indices) == numel
-            ), f"indices must be the same length as labels, got {len(indices)} and {numel}"
-        if weights is not None:
-            assert (
-                len(weights) == numel
-            ), f"weights must be the same length as labels, got {len(weights)} and {numel}"
-        if unique_ids is not None:
-            assert (
-                len(unique_ids) == numel
-            ), f"unique_ids must be the same length as labels, got {len(unique_ids)} and {numel}"
+
+        def _check_if_same_length(to_check, label):
+            """Simple utility method to throw an error if the length is not as expected."""
+            if to_check is not None and len(to_check) != numel:
+                raise ValueError(
+                    f"{label} must be the same length as `labels`, got {len(to_check)} and {numel}"
+                )
+
+        _check_if_same_length(features, "features")
+        _check_if_same_length(indices, "indices")
+        _check_if_same_length(weights, "weights")
+        _check_if_same_length(unique_ids, "unique_ids")
+        _check_if_same_length(mol_ids, "mol_ids")
 
         self.labels = labels
         if smiles is not None:
@@ -74,6 +69,7 @@ def __init__(
             )  # Avoid memory leaks with `num_workers > 0` by using numpy array
         self.weights = weights
         self.unique_ids = unique_ids
+        self.mol_ids = mol_ids
 
     def __len__(self):
         r"""
@@ -111,6 +107,9 @@ def __getitem__(self, idx):
         if self.unique_ids is not None:
             datum["unique_ids"] = self.unique_ids[idx]
 
+        if self.mol_ids is not None:
+            datum["mol_ids"] = self.mol_ids[idx]
+
         return datum
 
     def __getstate__(self):
@@ -122,6 +121,7 @@ def __getstate__(self):
         state["indices"] = self.indices
         state["weights"] = self.weights
         state["unique_ids"] = self.unique_ids
+        state["mol_ids"] = self.mol_ids
         return state
 
     def __setstate__(self, state: dict):
diff --git a/graphium/utils/spaces.py b/graphium/utils/spaces.py
index b88fb7d41..5ae052bb7 100644
--- a/graphium/utils/spaces.py
+++ b/graphium/utils/spaces.py
@@ -119,5 +119,6 @@
 DATAMODULE_DICT = {
     "GraphOGBDataModule": Datamodules.GraphOGBDataModule,
     "MultitaskFromSmilesDataModule": Datamodules.MultitaskFromSmilesDataModule,
+    "ADMETBenchmarkDataModule": Datamodules.ADMETBenchmarkDataModule,
     "FakeDataModule": Datamodules.FakeDataModule,
 }
diff --git a/notebooks/finetuning-on-tdc-admet-benchmark.ipynb b/notebooks/finetuning-on-tdc-admet-benchmark.ipynb
new file mode 100644
index 000000000..3c42c0349
--- /dev/null
+++ b/notebooks/finetuning-on-tdc-admet-benchmark.ipynb
@@ -0,0 +1,897 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "55ab828f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%load_ext autoreload\n",
+    "%autoreload 2"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "07e0ed51",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import yaml\n",
+    "import omegaconf\n",
+    "from datetime import datetime\n",
+    "\n",
+    "from typing import Union, List\n",
+    "from copy import deepcopy\n",
+    "from tdc.utils import retrieve_benchmark_names\n",
+    "\n",
+    "from graphium.config._loader import (\n",
+    "    load_datamodule,\n",
+    "    load_metrics,\n",
+    "    load_architecture,\n",
+    "    load_predictor,\n",
+    "    load_trainer,\n",
+    "    save_params_to_wandb,\n",
+    "    load_accelerator,\n",
+    "    load_yaml_config,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "fba4e4dc",
+   "metadata": {},
+   "source": [
+    "# Fine-tuning on the TDC ADMET benchmarking group\n",
+    "\n",
+    "[TDC](https://tdcommons.ai/) hosts a variety of ML-ready datasets and benchmarks for ML for drug discovery. The [TDC ADMET benchmarking group](https://tdcommons.ai/benchmark/admet_group/overview/) is a popular collection of benchmarks for evaluating new _foundation models_ (see e.g. [MolE](https://arxiv.org/abs/2211.02657)) due to the variety and relevance of the included tasks.\n",
+    "\n",
+    "The ADMET benchmarking group is integrated in `graphium` through the `ADMETBenchmarkDataModule` data-module. This notebook shows how to easily fine-tune and test a model using that data-module. \n",
+    "\n",
+    "<div style=\"background-color: #fff3cd; border-radius: 10px; border-color: #ffeeba; padding: 20px; margin: 20px 0;  color: #856404\">\n",
+    "    <b>NOTE:</b> This notebook is still <i>work in progress</i>. While the <b>fine-tuning logic is unfinished</b>, the notebook does demo how one could use the data-module to easily loop over each of the datasets in the benchmarking group and get the prescribed train-test split. Once the fine-tuning logic is finalized, we will finish this notebook and officially provide it as a tutorial within Graphium. \n",
+    "</div>"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "4d5af838",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# First, let's read the yaml configuration file\n",
+    "with open(\"../expts/configs/config_tdc_admet_demo.yaml\", \"r\") as file:\n",
+    "    config = yaml.load(file, Loader=yaml.FullLoader)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "7b2047a3",
+   "metadata": {},
+   "source": [
+    "## Get all TDC benchmark names"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "faa5d4b4",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "22"
+      ]
+     },
+     "execution_count": 4,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "benchmarks = retrieve_benchmark_names(\"admet_group\")\n",
+    "len(benchmarks)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f80de284",
+   "metadata": {},
+   "source": [
+    "While there is a total of 22, let's just use two for practicality sake: One regression and one classification task! "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "3ac3c111",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "benchmarks = [\"caco2_wang\", \"hia_hou\"]"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3643aa30",
+   "metadata": {},
+   "source": [
+    "## Initialize all training components per task\n",
+    "**NOTE**: Since we do not have fine-tuning logic, this for now just creates a new model. Ultimately, we will want to use fine-tuning code to evaluate how well the pre-trained model transfers to downstream tasks. "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "9538abfb",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def training_testing_loop(cfg):\n",
+    "    \"\"\"\n",
+    "    Simple loop to train a model from scratch and test it. \n",
+    "    \"\"\"\n",
+    "    \n",
+    "    # Initialize object from config\n",
+    "    cfg, accelerator_type = load_accelerator(cfg)\n",
+    "    datamodule = load_datamodule(cfg, accelerator_type)\n",
+    "    model_class, model_kwargs = load_architecture(cfg, in_dims=datamodule.in_dims)\n",
+    "    metrics = load_metrics(cfg)\n",
+    "    \n",
+    "    # Prepare data\n",
+    "    datamodule.prepare_data()\n",
+    "    \n",
+    "    # Initialize the predictor\n",
+    "    predictor = load_predictor(\n",
+    "        cfg, model_class, model_kwargs, metrics, accelerator_type, datamodule.task_norms\n",
+    "    )\n",
+    "    \n",
+    "    # Initialize the trainer\n",
+    "    date_time_suffix = datetime.now().strftime(\"%d.%m.%Y_%H.%M.%S\")\n",
+    "    trainer = load_trainer(cfg, \"tdc-admet\", accelerator_type, date_time_suffix)\n",
+    "        \n",
+    "    # Train\n",
+    "    predictor.set_max_nodes_edges_per_graph(datamodule, stages=[\"train\", \"val\"])\n",
+    "    trainer.fit(model=predictor, datamodule=datamodule)\n",
+    "    \n",
+    "    # Test\n",
+    "    predictor.set_max_nodes_edges_per_graph(datamodule, stages=[\"test\"])\n",
+    "    results = trainer.test(model=predictor, datamodule=datamodule)\n",
+    "    \n",
+    "    return results\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "1ee4586e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def filter_cfg_based_on_benchmark_name(config, names: Union[List[str], str]):\n",
+    "    \"\"\"\n",
+    "    Filter a base config for the full TDC ADMET benchmarking group to only \n",
+    "    have settings related to a subset of the endpoints\n",
+    "    \"\"\"\n",
+    "    \n",
+    "    if config[\"datamodule\"][\"module_type\"] != \"ADMETBenchmarkDataModule\":\n",
+    "        raise ValueError(\"You can only use this method for the `ADMETBenchmarkDataModule`\")\n",
+    "        \n",
+    "    if isinstance(names, str):\n",
+    "        names = [names]\n",
+    "    \n",
+    "    def _filter(d):\n",
+    "        return {k: v for k, v in d.items() if k in names}\n",
+    "         \n",
+    "    cfg = deepcopy(config)\n",
+    "    \n",
+    "    # Update the datamodule arguments\n",
+    "    cfg[\"datamodule\"][\"args\"][\"tdc_benchmark_names\"] = names\n",
+    "    \n",
+    "    # Filter the relevant config sections\n",
+    "    cfg[\"architecture\"][\"task_heads\"] = _filter(cfg[\"architecture\"][\"task_heads\"])\n",
+    "    cfg[\"predictor\"][\"metrics_on_progress_bar\"] = _filter(cfg[\"predictor\"][\"metrics_on_progress_bar\"])\n",
+    "    cfg[\"predictor\"][\"loss_fun\"] = _filter(cfg[\"predictor\"][\"loss_fun\"])\n",
+    "    cfg[\"metrics\"] = _filter(cfg[\"metrics\"])\n",
+    "    \n",
+    "    return cfg"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "96bc606a",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\u001b[32m2023-07-13 14:02:22.438\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36mgraphium.data.datamodule\u001b[0m:\u001b[36m__init__\u001b[0m:\u001b[36m2425\u001b[0m - \u001b[1mPreparing the TDC ADMET Benchmark Group splits for each of the 1 benchmarks.\u001b[0m\n",
+      "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33mcwognum\u001b[0m (\u001b[33mvalence-ood\u001b[0m). Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n",
+      "\u001b[34m\u001b[1mwandb\u001b[0m: \u001b[33mWARNING\u001b[0m Path logs/tdc-admet-demo/wandb/ wasn't writable, using system temp directory.\n",
+      "wandb: WARNING Path logs/tdc-admet-demo/wandb/ wasn't writable, using system temp directory\n"
+     ]
+    },
+    {
+     "data": {
+      "text/html": [
+       "Tracking run with wandb version 0.15.5"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "Run data is saved locally in <code>/tmp/wandb/run-20230713_140223-iablpj4z</code>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "Syncing run <strong><a href='https://wandb.ai/valence-ood/tdc_admet_demo/runs/iablpj4z' target=\"_blank\">tdc_admet_demo_13.07.2023_14.02.22</a></strong> to <a href='https://wandb.ai/valence-ood/tdc_admet_demo' target=\"_blank\">Weights & Biases</a> (<a href='https://wandb.me/run' target=\"_blank\">docs</a>)<br/>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       " View project at <a href='https://wandb.ai/valence-ood/tdc_admet_demo' target=\"_blank\">https://wandb.ai/valence-ood/tdc_admet_demo</a>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       " View run at <a href='https://wandb.ai/valence-ood/tdc_admet_demo/runs/iablpj4z' target=\"_blank\">https://wandb.ai/valence-ood/tdc_admet_demo/runs/iablpj4z</a>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "GPU available: True (cuda), used: True\n",
+      "TPU available: False, using: 0 TPU cores\n",
+      "IPU available: False, using: 0 IPUs\n",
+      "HPU available: False, using: 0 HPUs\n",
+      "\u001b[32m2023-07-13 14:02:28.065\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36mgraphium.data.datamodule\u001b[0m:\u001b[36msetup\u001b[0m:\u001b[36m1168\u001b[0m - \u001b[1m-------------------\n",
+      "MultitaskDataset\n",
+      "\tabout = training set\n",
+      "\tnum_graphs_total = 634\n",
+      "\tnum_nodes_total = 18351\n",
+      "\tmax_num_nodes_per_graph = 67\n",
+      "\tmin_num_nodes_per_graph = 2\n",
+      "\tstd_num_nodes_per_graph = 11.441048173903344\n",
+      "\tmean_num_nodes_per_graph = 28.944794952681388\n",
+      "\tnum_edges_total = 39466\n",
+      "\tmax_num_edges_per_graph = 144\n",
+      "\tmin_num_edges_per_graph = 2\n",
+      "\tstd_num_edges_per_graph = 25.32877270037731\n",
+      "\tmean_num_edges_per_graph = 62.24921135646688\n",
+      "-------------------\n",
+      "\u001b[0m\n",
+      "\u001b[32m2023-07-13 14:02:28.065\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36mgraphium.data.datamodule\u001b[0m:\u001b[36msetup\u001b[0m:\u001b[36m1169\u001b[0m - \u001b[1m-------------------\n",
+      "MultitaskDataset\n",
+      "\tabout = validation set\n",
+      "\tnum_graphs_total = 91\n",
+      "\tnum_nodes_total = 2791\n",
+      "\tmax_num_nodes_per_graph = 67\n",
+      "\tmin_num_nodes_per_graph = 13\n",
+      "\tstd_num_nodes_per_graph = 11.461406761225364\n",
+      "\tmean_num_nodes_per_graph = 30.67032967032967\n",
+      "\tnum_edges_total = 6034\n",
+      "\tmax_num_edges_per_graph = 148\n",
+      "\tmin_num_edges_per_graph = 28\n",
+      "\tstd_num_edges_per_graph = 25.619228830769817\n",
+      "\tmean_num_edges_per_graph = 66.3076923076923\n",
+      "-------------------\n",
+      "\u001b[0m\n",
+      "\u001b[32m2023-07-13 14:02:28.066\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36mgraphium.data.datamodule\u001b[0m:\u001b[36msetup\u001b[0m:\u001b[36m1186\u001b[0m - \u001b[1m-------------------\n",
+      "MultitaskDataset\n",
+      "\tabout = test set\n",
+      "\tnum_graphs_total = 181\n",
+      "\tnum_nodes_total = 5503\n",
+      "\tmax_num_nodes_per_graph = 68\n",
+      "\tmin_num_nodes_per_graph = 8\n",
+      "\tstd_num_nodes_per_graph = 9.244248028261115\n",
+      "\tmean_num_nodes_per_graph = 30.403314917127073\n",
+      "\tnum_edges_total = 11736\n",
+      "\tmax_num_edges_per_graph = 144\n",
+      "\tmin_num_edges_per_graph = 16\n",
+      "\tstd_num_edges_per_graph = 19.961337435561646\n",
+      "\tmean_num_edges_per_graph = 64.83977900552486\n",
+      "-------------------\n",
+      "\u001b[0m\n",
+      "\u001b[32m2023-07-13 14:02:28.066\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36mgraphium.data.datamodule\u001b[0m:\u001b[36mget_max_num_nodes_datamodule\u001b[0m:\u001b[36m556\u001b[0m - \u001b[1mMax num nodes being calcuated train\u001b[0m\n",
+      "\u001b[32m2023-07-13 14:02:28.067\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36mgraphium.data.datamodule\u001b[0m:\u001b[36mget_max_num_nodes_datamodule\u001b[0m:\u001b[36m565\u001b[0m - \u001b[1mMax num nodes being calcuated val\u001b[0m\n",
+      "\u001b[32m2023-07-13 14:02:28.067\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36mgraphium.data.datamodule\u001b[0m:\u001b[36mprepare_data\u001b[0m:\u001b[36m943\u001b[0m - \u001b[1mData is already prepared. Skipping the preparation\u001b[0m\n",
+      "You are using a CUDA device ('NVIDIA GeForce RTX 3070') that has Tensor Cores. To properly utilize them, you should set `torch.set_float32_matmul_precision('medium' | 'high')` which will trade-off precision for performance. For more details, read https://pytorch.org/docs/stable/generated/torch.set_float32_matmul_precision.html#torch.set_float32_matmul_precision\n",
+      "\u001b[32m2023-07-13 14:02:28.069\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36mgraphium.data.datamodule\u001b[0m:\u001b[36msetup\u001b[0m:\u001b[36m1168\u001b[0m - \u001b[1m-------------------\n",
+      "MultitaskDataset\n",
+      "\tabout = training set\n",
+      "\tnum_graphs_total = 634\n",
+      "\tnum_nodes_total = 18351\n",
+      "\tmax_num_nodes_per_graph = 67\n",
+      "\tmin_num_nodes_per_graph = 2\n",
+      "\tstd_num_nodes_per_graph = 11.441048173903344\n",
+      "\tmean_num_nodes_per_graph = 28.944794952681388\n",
+      "\tnum_edges_total = 39466\n",
+      "\tmax_num_edges_per_graph = 144\n",
+      "\tmin_num_edges_per_graph = 2\n",
+      "\tstd_num_edges_per_graph = 25.32877270037731\n",
+      "\tmean_num_edges_per_graph = 62.24921135646688\n",
+      "-------------------\n",
+      "\u001b[0m\n",
+      "\u001b[32m2023-07-13 14:02:28.069\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36mgraphium.data.datamodule\u001b[0m:\u001b[36msetup\u001b[0m:\u001b[36m1169\u001b[0m - \u001b[1m-------------------\n",
+      "MultitaskDataset\n",
+      "\tabout = validation set\n",
+      "\tnum_graphs_total = 91\n",
+      "\tnum_nodes_total = 2791\n",
+      "\tmax_num_nodes_per_graph = 67\n",
+      "\tmin_num_nodes_per_graph = 13\n",
+      "\tstd_num_nodes_per_graph = 11.461406761225364\n",
+      "\tmean_num_nodes_per_graph = 30.67032967032967\n",
+      "\tnum_edges_total = 6034\n",
+      "\tmax_num_edges_per_graph = 148\n",
+      "\tmin_num_edges_per_graph = 28\n",
+      "\tstd_num_edges_per_graph = 25.619228830769817\n",
+      "\tmean_num_edges_per_graph = 66.3076923076923\n",
+      "-------------------\n",
+      "\u001b[0m\n",
+      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]\n",
+      "\n",
+      "  | Name  | Type                      | Params\n",
+      "----------------------------------------------------\n",
+      "0 | model | FullGraphMultiTaskNetwork | 111 K \n",
+      "----------------------------------------------------\n",
+      "111 K     Trainable params\n",
+      "0         Non-trainable params\n",
+      "111 K     Total params\n",
+      "0.448     Total estimated model params size (MB)\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Sanity Checking: 0it [00:00, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "a99952d7025f418bae86c250887d7b60",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Training: 0it [00:00, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "`Trainer.fit` stopped: `max_epochs=10` reached.\n",
+      "\u001b[32m2023-07-13 14:02:49.408\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36mgraphium.data.datamodule\u001b[0m:\u001b[36msetup\u001b[0m:\u001b[36m1168\u001b[0m - \u001b[1m-------------------\n",
+      "MultitaskDataset\n",
+      "\tabout = training set\n",
+      "\tnum_graphs_total = 634\n",
+      "\tnum_nodes_total = 18351\n",
+      "\tmax_num_nodes_per_graph = 67\n",
+      "\tmin_num_nodes_per_graph = 2\n",
+      "\tstd_num_nodes_per_graph = 11.441048173903344\n",
+      "\tmean_num_nodes_per_graph = 28.944794952681388\n",
+      "\tnum_edges_total = 39466\n",
+      "\tmax_num_edges_per_graph = 144\n",
+      "\tmin_num_edges_per_graph = 2\n",
+      "\tstd_num_edges_per_graph = 25.32877270037731\n",
+      "\tmean_num_edges_per_graph = 62.24921135646688\n",
+      "-------------------\n",
+      "\u001b[0m\n",
+      "\u001b[32m2023-07-13 14:02:49.409\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36mgraphium.data.datamodule\u001b[0m:\u001b[36msetup\u001b[0m:\u001b[36m1169\u001b[0m - \u001b[1m-------------------\n",
+      "MultitaskDataset\n",
+      "\tabout = validation set\n",
+      "\tnum_graphs_total = 91\n",
+      "\tnum_nodes_total = 2791\n",
+      "\tmax_num_nodes_per_graph = 67\n",
+      "\tmin_num_nodes_per_graph = 13\n",
+      "\tstd_num_nodes_per_graph = 11.461406761225364\n",
+      "\tmean_num_nodes_per_graph = 30.67032967032967\n",
+      "\tnum_edges_total = 6034\n",
+      "\tmax_num_edges_per_graph = 148\n",
+      "\tmin_num_edges_per_graph = 28\n",
+      "\tstd_num_edges_per_graph = 25.619228830769817\n",
+      "\tmean_num_edges_per_graph = 66.3076923076923\n",
+      "-------------------\n",
+      "\u001b[0m\n",
+      "\u001b[32m2023-07-13 14:02:49.410\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36mgraphium.data.datamodule\u001b[0m:\u001b[36msetup\u001b[0m:\u001b[36m1186\u001b[0m - \u001b[1m-------------------\n",
+      "MultitaskDataset\n",
+      "\tabout = test set\n",
+      "\tnum_graphs_total = 181\n",
+      "\tnum_nodes_total = 5503\n",
+      "\tmax_num_nodes_per_graph = 68\n",
+      "\tmin_num_nodes_per_graph = 8\n",
+      "\tstd_num_nodes_per_graph = 9.244248028261115\n",
+      "\tmean_num_nodes_per_graph = 30.403314917127073\n",
+      "\tnum_edges_total = 11736\n",
+      "\tmax_num_edges_per_graph = 144\n",
+      "\tmin_num_edges_per_graph = 16\n",
+      "\tstd_num_edges_per_graph = 19.961337435561646\n",
+      "\tmean_num_edges_per_graph = 64.83977900552486\n",
+      "-------------------\n",
+      "\u001b[0m\n",
+      "\u001b[32m2023-07-13 14:02:49.410\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36mgraphium.data.datamodule\u001b[0m:\u001b[36mget_max_num_nodes_datamodule\u001b[0m:\u001b[36m574\u001b[0m - \u001b[1mMax num nodes being calcuated test\u001b[0m\n",
+      "\u001b[32m2023-07-13 14:02:49.411\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36mgraphium.data.datamodule\u001b[0m:\u001b[36mprepare_data\u001b[0m:\u001b[36m943\u001b[0m - \u001b[1mData is already prepared. Skipping the preparation\u001b[0m\n",
+      "You are using a CUDA device ('NVIDIA GeForce RTX 3070') that has Tensor Cores. To properly utilize them, you should set `torch.set_float32_matmul_precision('medium' | 'high')` which will trade-off precision for performance. For more details, read https://pytorch.org/docs/stable/generated/torch.set_float32_matmul_precision.html#torch.set_float32_matmul_precision\n",
+      "\u001b[32m2023-07-13 14:02:49.412\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36mgraphium.data.datamodule\u001b[0m:\u001b[36msetup\u001b[0m:\u001b[36m1186\u001b[0m - \u001b[1m-------------------\n",
+      "MultitaskDataset\n",
+      "\tabout = test set\n",
+      "\tnum_graphs_total = 181\n",
+      "\tnum_nodes_total = 5503\n",
+      "\tmax_num_nodes_per_graph = 68\n",
+      "\tmin_num_nodes_per_graph = 8\n",
+      "\tstd_num_nodes_per_graph = 9.244248028261115\n",
+      "\tmean_num_nodes_per_graph = 30.403314917127073\n",
+      "\tnum_edges_total = 11736\n",
+      "\tmax_num_edges_per_graph = 144\n",
+      "\tmin_num_edges_per_graph = 16\n",
+      "\tstd_num_edges_per_graph = 19.961337435561646\n",
+      "\tmean_num_edges_per_graph = 64.83977900552486\n",
+      "-------------------\n",
+      "\u001b[0m\n",
+      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "e92ded2754004662a47e2fb8cb8112be",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Testing: 0it [00:00, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
+       "┃<span style=\"font-weight: bold\">             Test metric             </span>┃<span style=\"font-weight: bold\">            DataLoader 0             </span>┃\n",
+       "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
+       "│<span style=\"color: #008080; text-decoration-color: #008080\">          gpu_allocated_GB           </span>│<span style=\"color: #800080; text-decoration-color: #800080\">        0.0019927024841308594        </span>│\n",
+       "│<span style=\"color: #008080; text-decoration-color: #008080\">    graph_caco2_wang/L1Loss/test     </span>│<span style=\"color: #800080; text-decoration-color: #800080\">          2.312683582305908          </span>│\n",
+       "│<span style=\"color: #008080; text-decoration-color: #008080\">     graph_caco2_wang/loss/test      </span>│<span style=\"color: #800080; text-decoration-color: #800080\">          2.312683582305908          </span>│\n",
+       "│<span style=\"color: #008080; text-decoration-color: #008080\">      graph_caco2_wang/mae/test      </span>│<span style=\"color: #800080; text-decoration-color: #800080\">          2.312683582305908          </span>│\n",
+       "│<span style=\"color: #008080; text-decoration-color: #008080\">   graph_caco2_wang/mean_pred/test   </span>│<span style=\"color: #800080; text-decoration-color: #800080\">         -3.0216422080993652         </span>│\n",
+       "│<span style=\"color: #008080; text-decoration-color: #008080\">  graph_caco2_wang/mean_target/test  </span>│<span style=\"color: #800080; text-decoration-color: #800080\">         -5.334325790405273          </span>│\n",
+       "│<span style=\"color: #008080; text-decoration-color: #008080\">  graph_caco2_wang/median_pred/test  </span>│<span style=\"color: #800080; text-decoration-color: #800080\">          -3.02557373046875          </span>│\n",
+       "│<span style=\"color: #008080; text-decoration-color: #008080\"> graph_caco2_wang/median_target/test </span>│<span style=\"color: #800080; text-decoration-color: #800080\">         -5.300000190734863          </span>│\n",
+       "│<span style=\"color: #008080; text-decoration-color: #008080\">    graph_caco2_wang/pearson/test    </span>│<span style=\"color: #800080; text-decoration-color: #800080\">         0.11455658078193665         </span>│\n",
+       "│<span style=\"color: #008080; text-decoration-color: #008080\">   graph_caco2_wang/r2_score/test    </span>│<span style=\"color: #800080; text-decoration-color: #800080\">         -11.43954086303711          </span>│\n",
+       "│<span style=\"color: #008080; text-decoration-color: #008080\">   graph_caco2_wang/spearman/test    </span>│<span style=\"color: #800080; text-decoration-color: #800080\">         0.2947588562965393          </span>│\n",
+       "│<span style=\"color: #008080; text-decoration-color: #008080\">   graph_caco2_wang/std_pred/test    </span>│<span style=\"color: #800080; text-decoration-color: #800080\">        0.014227418228983879         </span>│\n",
+       "│<span style=\"color: #008080; text-decoration-color: #008080\">  graph_caco2_wang/std_target/test   </span>│<span style=\"color: #800080; text-decoration-color: #800080\">         0.6855407357215881          </span>│\n",
+       "│<span style=\"color: #008080; text-decoration-color: #008080\">              loss/test              </span>│<span style=\"color: #800080; text-decoration-color: #800080\">          2.312683582305908          </span>│\n",
+       "└─────────────────────────────────────┴─────────────────────────────────────┘\n",
+       "</pre>\n"
+      ],
+      "text/plain": [
+       "┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
+       "┃\u001b[1m \u001b[0m\u001b[1m            Test metric            \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1m           DataLoader 0            \u001b[0m\u001b[1m \u001b[0m┃\n",
+       "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
+       "│\u001b[36m \u001b[0m\u001b[36m         gpu_allocated_GB          \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m       0.0019927024841308594       \u001b[0m\u001b[35m \u001b[0m│\n",
+       "│\u001b[36m \u001b[0m\u001b[36m   graph_caco2_wang/L1Loss/test    \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m         2.312683582305908         \u001b[0m\u001b[35m \u001b[0m│\n",
+       "│\u001b[36m \u001b[0m\u001b[36m    graph_caco2_wang/loss/test     \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m         2.312683582305908         \u001b[0m\u001b[35m \u001b[0m│\n",
+       "│\u001b[36m \u001b[0m\u001b[36m     graph_caco2_wang/mae/test     \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m         2.312683582305908         \u001b[0m\u001b[35m \u001b[0m│\n",
+       "│\u001b[36m \u001b[0m\u001b[36m  graph_caco2_wang/mean_pred/test  \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m        -3.0216422080993652        \u001b[0m\u001b[35m \u001b[0m│\n",
+       "│\u001b[36m \u001b[0m\u001b[36m graph_caco2_wang/mean_target/test \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m        -5.334325790405273         \u001b[0m\u001b[35m \u001b[0m│\n",
+       "│\u001b[36m \u001b[0m\u001b[36m graph_caco2_wang/median_pred/test \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m         -3.02557373046875         \u001b[0m\u001b[35m \u001b[0m│\n",
+       "│\u001b[36m \u001b[0m\u001b[36mgraph_caco2_wang/median_target/test\u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m        -5.300000190734863         \u001b[0m\u001b[35m \u001b[0m│\n",
+       "│\u001b[36m \u001b[0m\u001b[36m   graph_caco2_wang/pearson/test   \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m        0.11455658078193665        \u001b[0m\u001b[35m \u001b[0m│\n",
+       "│\u001b[36m \u001b[0m\u001b[36m  graph_caco2_wang/r2_score/test   \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m        -11.43954086303711         \u001b[0m\u001b[35m \u001b[0m│\n",
+       "│\u001b[36m \u001b[0m\u001b[36m  graph_caco2_wang/spearman/test   \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m        0.2947588562965393         \u001b[0m\u001b[35m \u001b[0m│\n",
+       "│\u001b[36m \u001b[0m\u001b[36m  graph_caco2_wang/std_pred/test   \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m       0.014227418228983879        \u001b[0m\u001b[35m \u001b[0m│\n",
+       "│\u001b[36m \u001b[0m\u001b[36m graph_caco2_wang/std_target/test  \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m        0.6855407357215881         \u001b[0m\u001b[35m \u001b[0m│\n",
+       "│\u001b[36m \u001b[0m\u001b[36m             loss/test             \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m         2.312683582305908         \u001b[0m\u001b[35m \u001b[0m│\n",
+       "└─────────────────────────────────────┴─────────────────────────────────────┘\n"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\u001b[32m2023-07-13 14:02:50.286\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36mgraphium.data.datamodule\u001b[0m:\u001b[36m__init__\u001b[0m:\u001b[36m2425\u001b[0m - \u001b[1mPreparing the TDC ADMET Benchmark Group splits for each of the 1 benchmarks.\u001b[0m\n",
+      "GPU available: True (cuda), used: True\n",
+      "TPU available: False, using: 0 TPU cores\n",
+      "IPU available: False, using: 0 IPUs\n",
+      "HPU available: False, using: 0 HPUs\n",
+      "\u001b[32m2023-07-13 14:02:50.421\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36mgraphium.data.datamodule\u001b[0m:\u001b[36msetup\u001b[0m:\u001b[36m1168\u001b[0m - \u001b[1m-------------------\n",
+      "MultitaskDataset\n",
+      "\tabout = training set\n",
+      "\tnum_graphs_total = 403\n",
+      "\tnum_nodes_total = 9065\n",
+      "\tmax_num_nodes_per_graph = 101\n",
+      "\tmin_num_nodes_per_graph = 7\n",
+      "\tstd_num_nodes_per_graph = 8.379156239363269\n",
+      "\tmean_num_nodes_per_graph = 22.49379652605459\n",
+      "\tnum_edges_total = 19420\n",
+      "\tmax_num_edges_per_graph = 220\n",
+      "\tmin_num_edges_per_graph = 14\n",
+      "\tstd_num_edges_per_graph = 18.678530787820403\n",
+      "\tmean_num_edges_per_graph = 48.188585607940446\n",
+      "-------------------\n",
+      "\u001b[0m\n",
+      "\u001b[32m2023-07-13 14:02:50.422\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36mgraphium.data.datamodule\u001b[0m:\u001b[36msetup\u001b[0m:\u001b[36m1169\u001b[0m - \u001b[1m-------------------\n",
+      "MultitaskDataset\n",
+      "\tabout = validation set\n",
+      "\tnum_graphs_total = 58\n",
+      "\tnum_nodes_total = 1431\n",
+      "\tmax_num_nodes_per_graph = 67\n",
+      "\tmin_num_nodes_per_graph = 9\n",
+      "\tstd_num_nodes_per_graph = 9.610317607573146\n",
+      "\tmean_num_nodes_per_graph = 24.67241379310345\n",
+      "\tnum_edges_total = 3102\n",
+      "\tmax_num_edges_per_graph = 144\n",
+      "\tmin_num_edges_per_graph = 18\n",
+      "\tstd_num_edges_per_graph = 21.938976007372833\n",
+      "\tmean_num_edges_per_graph = 53.48275862068966\n",
+      "-------------------\n",
+      "\u001b[0m\n",
+      "\u001b[32m2023-07-13 14:02:50.423\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36mgraphium.data.datamodule\u001b[0m:\u001b[36msetup\u001b[0m:\u001b[36m1186\u001b[0m - \u001b[1m-------------------\n",
+      "MultitaskDataset\n",
+      "\tabout = test set\n",
+      "\tnum_graphs_total = 117\n",
+      "\tnum_nodes_total = 2716\n",
+      "\tmax_num_nodes_per_graph = 66\n",
+      "\tmin_num_nodes_per_graph = 3\n",
+      "\tstd_num_nodes_per_graph = 9.966464331904731\n",
+      "\tmean_num_nodes_per_graph = 23.213675213675213\n",
+      "\tnum_edges_total = 5790\n",
+      "\tmax_num_edges_per_graph = 136\n",
+      "\tmin_num_edges_per_graph = 4\n",
+      "\tstd_num_edges_per_graph = 22.25440142180863\n",
+      "\tmean_num_edges_per_graph = 49.48717948717949\n",
+      "-------------------\n",
+      "\u001b[0m\n",
+      "\u001b[32m2023-07-13 14:02:50.423\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36mgraphium.data.datamodule\u001b[0m:\u001b[36mget_max_num_nodes_datamodule\u001b[0m:\u001b[36m556\u001b[0m - \u001b[1mMax num nodes being calcuated train\u001b[0m\n",
+      "\u001b[32m2023-07-13 14:02:50.423\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36mgraphium.data.datamodule\u001b[0m:\u001b[36mget_max_num_nodes_datamodule\u001b[0m:\u001b[36m565\u001b[0m - \u001b[1mMax num nodes being calcuated val\u001b[0m\n",
+      "\u001b[32m2023-07-13 14:02:50.424\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36mgraphium.data.datamodule\u001b[0m:\u001b[36mprepare_data\u001b[0m:\u001b[36m943\u001b[0m - \u001b[1mData is already prepared. Skipping the preparation\u001b[0m\n",
+      "You are using a CUDA device ('NVIDIA GeForce RTX 3070') that has Tensor Cores. To properly utilize them, you should set `torch.set_float32_matmul_precision('medium' | 'high')` which will trade-off precision for performance. For more details, read https://pytorch.org/docs/stable/generated/torch.set_float32_matmul_precision.html#torch.set_float32_matmul_precision\n",
+      "\u001b[32m2023-07-13 14:02:50.426\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36mgraphium.data.datamodule\u001b[0m:\u001b[36msetup\u001b[0m:\u001b[36m1168\u001b[0m - \u001b[1m-------------------\n",
+      "MultitaskDataset\n",
+      "\tabout = training set\n",
+      "\tnum_graphs_total = 403\n",
+      "\tnum_nodes_total = 9065\n",
+      "\tmax_num_nodes_per_graph = 101\n",
+      "\tmin_num_nodes_per_graph = 7\n",
+      "\tstd_num_nodes_per_graph = 8.379156239363269\n",
+      "\tmean_num_nodes_per_graph = 22.49379652605459\n",
+      "\tnum_edges_total = 19420\n",
+      "\tmax_num_edges_per_graph = 220\n",
+      "\tmin_num_edges_per_graph = 14\n",
+      "\tstd_num_edges_per_graph = 18.678530787820403\n",
+      "\tmean_num_edges_per_graph = 48.188585607940446\n",
+      "-------------------\n",
+      "\u001b[0m\n",
+      "\u001b[32m2023-07-13 14:02:50.426\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36mgraphium.data.datamodule\u001b[0m:\u001b[36msetup\u001b[0m:\u001b[36m1169\u001b[0m - \u001b[1m-------------------\n",
+      "MultitaskDataset\n",
+      "\tabout = validation set\n",
+      "\tnum_graphs_total = 58\n",
+      "\tnum_nodes_total = 1431\n",
+      "\tmax_num_nodes_per_graph = 67\n",
+      "\tmin_num_nodes_per_graph = 9\n",
+      "\tstd_num_nodes_per_graph = 9.610317607573146\n",
+      "\tmean_num_nodes_per_graph = 24.67241379310345\n",
+      "\tnum_edges_total = 3102\n",
+      "\tmax_num_edges_per_graph = 144\n",
+      "\tmin_num_edges_per_graph = 18\n",
+      "\tstd_num_edges_per_graph = 21.938976007372833\n",
+      "\tmean_num_edges_per_graph = 53.48275862068966\n",
+      "-------------------\n",
+      "\u001b[0m\n",
+      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]\n",
+      "\n",
+      "  | Name  | Type                      | Params\n",
+      "----------------------------------------------------\n",
+      "0 | model | FullGraphMultiTaskNetwork | 111 K \n",
+      "----------------------------------------------------\n",
+      "111 K     Trainable params\n",
+      "0         Non-trainable params\n",
+      "111 K     Total params\n",
+      "0.448     Total estimated model params size (MB)\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Sanity Checking: 0it [00:00, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\u001b[32m2023-07-13 14:02:51.236\u001b[0m | \u001b[33m\u001b[1mWARNING \u001b[0m | \u001b[36mgraphium.trainer.predictor_summaries\u001b[0m:\u001b[36mget_metrics_logs\u001b[0m:\u001b[36m273\u001b[0m - \u001b[33m\u001b[1mError for metric graph_hia_hou/mcc/val. NaN is returned. Exception: stack expects a non-empty TensorList\u001b[0m\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "9ab41b110acd4f3f8977249dc715e49d",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Training: 0it [00:00, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\u001b[32m2023-07-13 14:02:52.078\u001b[0m | \u001b[33m\u001b[1mWARNING \u001b[0m | \u001b[36mgraphium.trainer.predictor_summaries\u001b[0m:\u001b[36mget_metrics_logs\u001b[0m:\u001b[36m273\u001b[0m - \u001b[33m\u001b[1mError for metric graph_hia_hou/mcc/train. NaN is returned. Exception: stack expects a non-empty TensorList\u001b[0m\n",
+      "`Trainer.fit` stopped: `max_epochs=10` reached.\n",
+      "\u001b[32m2023-07-13 14:03:06.289\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36mgraphium.data.datamodule\u001b[0m:\u001b[36msetup\u001b[0m:\u001b[36m1168\u001b[0m - \u001b[1m-------------------\n",
+      "MultitaskDataset\n",
+      "\tabout = training set\n",
+      "\tnum_graphs_total = 403\n",
+      "\tnum_nodes_total = 9065\n",
+      "\tmax_num_nodes_per_graph = 101\n",
+      "\tmin_num_nodes_per_graph = 7\n",
+      "\tstd_num_nodes_per_graph = 8.379156239363269\n",
+      "\tmean_num_nodes_per_graph = 22.49379652605459\n",
+      "\tnum_edges_total = 19420\n",
+      "\tmax_num_edges_per_graph = 220\n",
+      "\tmin_num_edges_per_graph = 14\n",
+      "\tstd_num_edges_per_graph = 18.678530787820403\n",
+      "\tmean_num_edges_per_graph = 48.188585607940446\n",
+      "-------------------\n",
+      "\u001b[0m\n",
+      "\u001b[32m2023-07-13 14:03:06.290\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36mgraphium.data.datamodule\u001b[0m:\u001b[36msetup\u001b[0m:\u001b[36m1169\u001b[0m - \u001b[1m-------------------\n",
+      "MultitaskDataset\n",
+      "\tabout = validation set\n",
+      "\tnum_graphs_total = 58\n",
+      "\tnum_nodes_total = 1431\n",
+      "\tmax_num_nodes_per_graph = 67\n",
+      "\tmin_num_nodes_per_graph = 9\n",
+      "\tstd_num_nodes_per_graph = 9.610317607573146\n",
+      "\tmean_num_nodes_per_graph = 24.67241379310345\n",
+      "\tnum_edges_total = 3102\n",
+      "\tmax_num_edges_per_graph = 144\n",
+      "\tmin_num_edges_per_graph = 18\n",
+      "\tstd_num_edges_per_graph = 21.938976007372833\n",
+      "\tmean_num_edges_per_graph = 53.48275862068966\n",
+      "-------------------\n",
+      "\u001b[0m\n",
+      "\u001b[32m2023-07-13 14:03:06.290\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36mgraphium.data.datamodule\u001b[0m:\u001b[36msetup\u001b[0m:\u001b[36m1186\u001b[0m - \u001b[1m-------------------\n",
+      "MultitaskDataset\n",
+      "\tabout = test set\n",
+      "\tnum_graphs_total = 117\n",
+      "\tnum_nodes_total = 2716\n",
+      "\tmax_num_nodes_per_graph = 66\n",
+      "\tmin_num_nodes_per_graph = 3\n",
+      "\tstd_num_nodes_per_graph = 9.966464331904731\n",
+      "\tmean_num_nodes_per_graph = 23.213675213675213\n",
+      "\tnum_edges_total = 5790\n",
+      "\tmax_num_edges_per_graph = 136\n",
+      "\tmin_num_edges_per_graph = 4\n",
+      "\tstd_num_edges_per_graph = 22.25440142180863\n",
+      "\tmean_num_edges_per_graph = 49.48717948717949\n",
+      "-------------------\n",
+      "\u001b[0m\n",
+      "\u001b[32m2023-07-13 14:03:06.291\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36mgraphium.data.datamodule\u001b[0m:\u001b[36mget_max_num_nodes_datamodule\u001b[0m:\u001b[36m574\u001b[0m - \u001b[1mMax num nodes being calcuated test\u001b[0m\n",
+      "\u001b[32m2023-07-13 14:03:06.291\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36mgraphium.data.datamodule\u001b[0m:\u001b[36mprepare_data\u001b[0m:\u001b[36m943\u001b[0m - \u001b[1mData is already prepared. Skipping the preparation\u001b[0m\n",
+      "You are using a CUDA device ('NVIDIA GeForce RTX 3070') that has Tensor Cores. To properly utilize them, you should set `torch.set_float32_matmul_precision('medium' | 'high')` which will trade-off precision for performance. For more details, read https://pytorch.org/docs/stable/generated/torch.set_float32_matmul_precision.html#torch.set_float32_matmul_precision\n",
+      "\u001b[32m2023-07-13 14:03:06.292\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36mgraphium.data.datamodule\u001b[0m:\u001b[36msetup\u001b[0m:\u001b[36m1186\u001b[0m - \u001b[1m-------------------\n",
+      "MultitaskDataset\n",
+      "\tabout = test set\n",
+      "\tnum_graphs_total = 117\n",
+      "\tnum_nodes_total = 2716\n",
+      "\tmax_num_nodes_per_graph = 66\n",
+      "\tmin_num_nodes_per_graph = 3\n",
+      "\tstd_num_nodes_per_graph = 9.966464331904731\n",
+      "\tmean_num_nodes_per_graph = 23.213675213675213\n",
+      "\tnum_edges_total = 5790\n",
+      "\tmax_num_edges_per_graph = 136\n",
+      "\tmin_num_edges_per_graph = 4\n",
+      "\tstd_num_edges_per_graph = 22.25440142180863\n",
+      "\tmean_num_edges_per_graph = 49.48717948717949\n",
+      "-------------------\n",
+      "\u001b[0m\n",
+      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "47a42b8385de4cd484d3aee1ee77a8ca",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Testing: 0it [00:00, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\u001b[32m2023-07-13 14:03:07.087\u001b[0m | \u001b[33m\u001b[1mWARNING \u001b[0m | \u001b[36mgraphium.trainer.predictor_summaries\u001b[0m:\u001b[36mget_metrics_logs\u001b[0m:\u001b[36m273\u001b[0m - \u001b[33m\u001b[1mError for metric graph_hia_hou/mcc/test. NaN is returned. Exception: stack expects a non-empty TensorList\u001b[0m\n"
+     ]
+    },
+    {
+     "data": {
+      "text/html": [
+       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
+       "┃<span style=\"font-weight: bold\">           Test metric            </span>┃<span style=\"font-weight: bold\">           DataLoader 0           </span>┃\n",
+       "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
+       "│<span style=\"color: #008080; text-decoration-color: #008080\">         gpu_allocated_GB         </span>│<span style=\"color: #800080; text-decoration-color: #800080\">      0.0018811225891113281       </span>│\n",
+       "│<span style=\"color: #008080; text-decoration-color: #008080\">    graph_hia_hou/BCELoss/test    </span>│<span style=\"color: #800080; text-decoration-color: #800080\">        0.5804430842399597        </span>│\n",
+       "│<span style=\"color: #008080; text-decoration-color: #008080\">   graph_hia_hou/accuracy/test    </span>│<span style=\"color: #800080; text-decoration-color: #800080\">        0.7692307829856873        </span>│\n",
+       "│<span style=\"color: #008080; text-decoration-color: #008080\">     graph_hia_hou/auprc/test     </span>│<span style=\"color: #800080; text-decoration-color: #800080\">        0.8519759774208069        </span>│\n",
+       "│<span style=\"color: #008080; text-decoration-color: #008080\">     graph_hia_hou/auroc/test     </span>│<span style=\"color: #800080; text-decoration-color: #800080\">        0.7008230686187744        </span>│\n",
+       "│<span style=\"color: #008080; text-decoration-color: #008080\">     graph_hia_hou/loss/test      </span>│<span style=\"color: #800080; text-decoration-color: #800080\">        0.5804430842399597        </span>│\n",
+       "│<span style=\"color: #008080; text-decoration-color: #008080\">      graph_hia_hou/mcc/test      </span>│<span style=\"color: #800080; text-decoration-color: #800080\">               nan                </span>│\n",
+       "│<span style=\"color: #008080; text-decoration-color: #008080\">   graph_hia_hou/mean_pred/test   </span>│<span style=\"color: #800080; text-decoration-color: #800080\">        0.8799146413803101        </span>│\n",
+       "│<span style=\"color: #008080; text-decoration-color: #008080\">  graph_hia_hou/mean_target/test  </span>│<span style=\"color: #800080; text-decoration-color: #800080\">        0.7692307829856873        </span>│\n",
+       "│<span style=\"color: #008080; text-decoration-color: #008080\">  graph_hia_hou/median_pred/test  </span>│<span style=\"color: #800080; text-decoration-color: #800080\">        0.8844738006591797        </span>│\n",
+       "│<span style=\"color: #008080; text-decoration-color: #008080\"> graph_hia_hou/median_target/test </span>│<span style=\"color: #800080; text-decoration-color: #800080\">               1.0                </span>│\n",
+       "│<span style=\"color: #008080; text-decoration-color: #008080\">   graph_hia_hou/std_pred/test    </span>│<span style=\"color: #800080; text-decoration-color: #800080\">       0.010796700604259968       </span>│\n",
+       "│<span style=\"color: #008080; text-decoration-color: #008080\">  graph_hia_hou/std_target/test   </span>│<span style=\"color: #800080; text-decoration-color: #800080\">       0.42313718795776367        </span>│\n",
+       "│<span style=\"color: #008080; text-decoration-color: #008080\">            loss/test             </span>│<span style=\"color: #800080; text-decoration-color: #800080\">        0.5804430842399597        </span>│\n",
+       "└──────────────────────────────────┴──────────────────────────────────┘\n",
+       "</pre>\n"
+      ],
+      "text/plain": [
+       "┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓\n",
+       "┃\u001b[1m \u001b[0m\u001b[1m          Test metric           \u001b[0m\u001b[1m \u001b[0m┃\u001b[1m \u001b[0m\u001b[1m          DataLoader 0          \u001b[0m\u001b[1m \u001b[0m┃\n",
+       "┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩\n",
+       "│\u001b[36m \u001b[0m\u001b[36m        gpu_allocated_GB        \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m     0.0018811225891113281      \u001b[0m\u001b[35m \u001b[0m│\n",
+       "│\u001b[36m \u001b[0m\u001b[36m   graph_hia_hou/BCELoss/test   \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m       0.5804430842399597       \u001b[0m\u001b[35m \u001b[0m│\n",
+       "│\u001b[36m \u001b[0m\u001b[36m  graph_hia_hou/accuracy/test   \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m       0.7692307829856873       \u001b[0m\u001b[35m \u001b[0m│\n",
+       "│\u001b[36m \u001b[0m\u001b[36m    graph_hia_hou/auprc/test    \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m       0.8519759774208069       \u001b[0m\u001b[35m \u001b[0m│\n",
+       "│\u001b[36m \u001b[0m\u001b[36m    graph_hia_hou/auroc/test    \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m       0.7008230686187744       \u001b[0m\u001b[35m \u001b[0m│\n",
+       "│\u001b[36m \u001b[0m\u001b[36m    graph_hia_hou/loss/test     \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m       0.5804430842399597       \u001b[0m\u001b[35m \u001b[0m│\n",
+       "│\u001b[36m \u001b[0m\u001b[36m     graph_hia_hou/mcc/test     \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m              nan               \u001b[0m\u001b[35m \u001b[0m│\n",
+       "│\u001b[36m \u001b[0m\u001b[36m  graph_hia_hou/mean_pred/test  \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m       0.8799146413803101       \u001b[0m\u001b[35m \u001b[0m│\n",
+       "│\u001b[36m \u001b[0m\u001b[36m graph_hia_hou/mean_target/test \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m       0.7692307829856873       \u001b[0m\u001b[35m \u001b[0m│\n",
+       "│\u001b[36m \u001b[0m\u001b[36m graph_hia_hou/median_pred/test \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m       0.8844738006591797       \u001b[0m\u001b[35m \u001b[0m│\n",
+       "│\u001b[36m \u001b[0m\u001b[36mgraph_hia_hou/median_target/test\u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m              1.0               \u001b[0m\u001b[35m \u001b[0m│\n",
+       "│\u001b[36m \u001b[0m\u001b[36m  graph_hia_hou/std_pred/test   \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m      0.010796700604259968      \u001b[0m\u001b[35m \u001b[0m│\n",
+       "│\u001b[36m \u001b[0m\u001b[36m graph_hia_hou/std_target/test  \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m      0.42313718795776367       \u001b[0m\u001b[35m \u001b[0m│\n",
+       "│\u001b[36m \u001b[0m\u001b[36m           loss/test            \u001b[0m\u001b[36m \u001b[0m│\u001b[35m \u001b[0m\u001b[35m       0.5804430842399597       \u001b[0m\u001b[35m \u001b[0m│\n",
+       "└──────────────────────────────────┴──────────────────────────────────┘\n"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "results = {}\n",
+    "\n",
+    "for name in benchmarks: \n",
+    "    \n",
+    "    # Run the training-testing loop\n",
+    "    cfg = filter_cfg_based_on_benchmark_name(config, name)\n",
+    "    benchmark_results = training_testing_loop(cfg)\n",
+    "    \n",
+    "    # Extract the main metric from the config\n",
+    "    metric = cfg[\"predictor\"][\"metrics_on_progress_bar\"][name][0]\n",
+    "    key = f\"graph_{name}/{metric}/test\"\n",
+    "    results[f\"{name}/{metric}\"] = benchmark_results[0][key]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "c2cea93b",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "caco2_wang/mae: 2.312683582305908\n",
+      "hia_hou/auroc: 0.7008230686187744\n",
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "print(omegaconf.OmegaConf.to_yaml(results))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "2150be90",
+   "metadata": {},
+   "source": [
+    "The End. "
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.12"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}


Test patch:
diff --git a/tests/test_multitask_datamodule.py b/tests/test_multitask_datamodule.py
index 965cc0a02..09dc53f41 100644
--- a/tests/test_multitask_datamodule.py
+++ b/tests/test_multitask_datamodule.py
@@ -1,4 +1,8 @@
+import shutil
+import tempfile
 import unittest as ut
+
+import pytest
 from omegaconf import OmegaConf
 import pandas as pd
 import numpy as np
@@ -6,6 +10,14 @@
 
 
 class Test_Multitask_DataModule(ut.TestCase):
+    def setUp(self):
+        # Create a temporary directory
+        self.tmp_test_dir = tempfile.mkdtemp()
+
+    def tearDown(self):
+        # Remove the directory after the test
+        shutil.rmtree(self.tmp_test_dir)
+
     def test_multitask_fromsmiles_dm(
         self,
     ):  # TODO: I think we can remove this as it tests tiny_zinc which only contain graph level labels
@@ -340,6 +352,56 @@ def test_non_graph_level_extract_labels_missing_cols(self):
                         non_missing_col = label_cols[1]
                         assert output[idx].shape[0] == len(df[non_missing_col][idx])
 
+    def test_tdc_admet_benchmark_data_module(self):
+        """
+        Verifies that the ADMET-specific subclass of the MultiTaskDataModule works.
+        Checks if all main endpoints can be run and if the split is correct.
+        """
+
+        try:
+            from tdc.benchmark_group import admet_group
+            from tdc.utils import retrieve_benchmark_names
+        except ImportError:
+            self.skipTest("PyTDC needs to be installed to run this test. Use `pip install PyTDC`.")
+            raise
+
+        # Make sure we can initialize the module and run the main endpoints
+        data_module = graphium.data.ADMETBenchmarkDataModule()
+        data_module.prepare_data()
+        data_module.setup()
+
+        for dl in [
+            data_module.train_dataloader(),
+            data_module.val_dataloader(),
+            data_module.test_dataloader(),
+        ]:
+            batch = next(iter(dl))
+            assert set(batch.keys()) == {"labels", "features"}
+
+        # # Validate the split
+        group = admet_group(path=self.tmp_test_dir)
+        benchmark_names = retrieve_benchmark_names("admet_group")
+
+        # For each of the endpoints...
+        for name in benchmark_names:
+            # Get the split from the benchmark group (ground truth)
+            benchmark = group.get(name)
+            train, val = group.get_train_valid_split(0, name)
+            test = benchmark["test"]
+
+            # Get the split from the data module
+            params = data_module._get_task_specific_arguments(name, 0, self.tmp_test_dir)
+            split = pd.read_csv(params.splits_path)
+            data = params.df
+
+            # Check that the split is the same
+            for ground_truth, label in [(train, "train"), (val, "val"), (test, "test")]:
+                y_true = ground_truth["Y"].values
+                y_module = data.loc[split[label].dropna(), "Y"].values
+
+                assert len(y_true) == len(y_module)
+                assert np.allclose(np.sort(y_true), np.sort(y_module))
+
 
 if __name__ == "__main__":
     ut.main()