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

Issue:
Support for multi-GPUs
Code crashes when using multiple GPUs, with some of the tensors being on the CPU, thus causing an error. To investigate more.

- Comment out every call to `.device`
- 


Discussion / hints:
I investigated the issue and it seems to be due to an incompatibility between DGL and Pytorch-Lightning. I know that DGL is working on it, but I don't think there's an easy fix as of now. I raised an issue on the DGL repository and we'll wait for updates.
See raised issue here: https://github.com/dmlc/dgl/issues/3053
According to the dgl ticket, `ddp` should work ( not `dp`) but we must use their custom `GraphDataLoader`: https://github.com/dmlc/dgl/issues/3053#issuecomment-869283620

> For the former, currently DGL only supports DistributedDataParallel, not DataParallel. So I guess you will need to try GraphDataLoader (with use_ddp option) and the ddp accelerator. Please let me know if you cannot get ddp working either (and preferably with your code so I can reproduce it).
Lots of work from @graphriela , but without success. The issue will be on hold until DGL or PL come up with fixes themselves.
Will close and re investigate when migrating to pytorch geometric

Fix patch:
diff --git a/.gitignore b/.gitignore
index 461406a83..146dda239 100644
--- a/.gitignore
+++ b/.gitignore
@@ -22,6 +22,7 @@ wandb/
 out/
 datacache/
 tests/temp_cache*
+predictions/
 
 # Data and predictions
 goli/data/ZINC_bench_gnn/
@@ -31,7 +32,13 @@ goli/data/cache/
 goli/data/b3lyp/
 goli/data/PCQM4Mv2/
 goli/data/PCQM4M/
-predictions/
+goli/data/neurips2023/small-dataset/
+goli/data/neurips2023/large-dataset/
+goli/data/neurips2023/dummy-dataset/
+goli/data/make_data_splits/*.csv*
+goli/data/make_data_splits/*.pt*
+goli/data/make_data_splits/*.parquet*
+
 
 # Others
 expts_untracked/
diff --git a/expts/configs/config_gps_10M_pcqm4m.yaml b/expts/configs/config_gps_10M_pcqm4m.yaml
index c7a7525b1..5ec9cef46 100644
--- a/expts/configs/config_gps_10M_pcqm4m.yaml
+++ b/expts/configs/config_gps_10M_pcqm4m.yaml
@@ -64,6 +64,7 @@ datamodule:
     task_specific_args:   # To be replaced by a new class "DatasetParams"
       homolumo:
         df: null
+        task_level: "graph"
         df_path: ~/scratch/data/goli/data/PCQM4M/pcqm4mv2-20k.csv
         # wget https://storage.googleapis.com/datasets-public-research/PCQM4M/cxsmiles/pcqm4mv2-20k.csv
         # or set path as https://storage.googleapis.com/datasets-public-research/PCQM4M/cxsmiles/pcqm4mv2-20k.csv directly
diff --git a/expts/configs/config_gps_10M_pcqm4m_mod.yaml b/expts/configs/config_gps_10M_pcqm4m_mod.yaml
new file mode 100644
index 000000000..2f8b509dc
--- /dev/null
+++ b/expts/configs/config_gps_10M_pcqm4m_mod.yaml
@@ -0,0 +1,289 @@
+# Testing the mpnn only model with the PCQMv2 dataset on IPU.
+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
+
+datamodule:
+  module_type: "MultitaskFromSmilesDataModule"
+  # module_type: "FakeDataModule"  # Option to use generated data
+  args: # Matches that in the test_multitask_datamodule.py case.
+    task_specific_args:   # To be replaced by a new class "DatasetParams"
+      homolumo:
+        df: null
+        task_level: "graph"
+        df_path: ~/scratch/data/goli/data/PCQM4M/pcqm4mv2-20k.csv
+        # wget https://storage.googleapis.com/datasets-public-research/PCQM4M/cxsmiles/pcqm4mv2-20k.csv
+        # or set path as https://storage.googleapis.com/datasets-public-research/PCQM4M/cxsmiles/pcqm4mv2-20k.csv directly
+        smiles_col: "cxsmiles"
+        label_cols: ["homo_lumo_gap"]
+        # sample_size: 30000 # use sample_size for test
+        # splits_path: goli/data/PCQM4Mv2/split_dict.pt  # Download with `wget https://storage.googleapis.com/datasets-public-research/PCQM4M/cxsmiles/split_dict.pt`
+        split_val: 0.1
+        split_test: 0.1
+
+    # Featurization
+    prepare_dict_or_graph: pyg:graph
+    featurization_n_jobs: 30
+    featurization_progress: True
+    featurization_backend: "loky"
+    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)']
+      mask_nan: 0
+      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]
+      conformer_property_list: [positions_3d]
+      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:
+          node_laplacian_eigvec:
+            pos_type: laplacian_eigvec
+            pos_level: node
+            num_pos: 8
+            normalization: "none"
+            disconnected_comp: True
+          node_laplacian_eigval:
+            pos_type: laplacian_eigval
+            pos_level: node
+            num_pos: 8
+            normalization: "none"
+            disconnected_comp: True
+          rw_return_probs:
+            pos_type: rw_return_probs
+            pos_level: node
+            ksteps: [4, 8]
+          nodepair_rw_transition_probs:
+            pos_type: rw_transition_probs
+            pos_level: edge
+            ksteps: [2, 4]
+          nodepair_rw_return_probs:
+            pos_type: rw_return_probs
+            pos_level: nodepair
+            ksteps: [4]
+          electrostatic:
+            pos_type: electrostatic
+            pos_level: node
+          edge_commute:
+            pos_type: commute
+            pos_level: edge
+          nodepair_graphormer:
+            pos_type: graphormer
+            pos_level: nodepair
+      
+    # Data handling-related
+    batch_size_training: 64
+    batch_size_inference: 16
+    # cache_data_path: .
+    num_workers: 0 # -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.
+    featurization_backend: "loky"
+
+    # 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
+  pre_nn:   # Set as null to avoid a pre-nn network
+    out_dim: 32
+    hidden_dims: 64
+    depth: 2
+    activation: relu
+    last_activation: none
+    dropout: &dropout 0.1
+    normalization: &normalization layer_norm
+    last_normalization: *normalization
+    residual_type: none
+
+  pre_nn_edges:   # Set as null to avoid a pre-nn network
+    out_dim: 16
+    hidden_dims: 32
+    depth: 2
+    activation: relu
+    last_activation: none
+    dropout: *dropout
+    normalization: *normalization
+    last_normalization: *normalization
+    residual_type: none
+
+  pe_encoders:
+    out_dim: &pe_out_dim 32
+    edge_out_dim: &edge_pe_out_dim 16
+    pool: "sum" #"mean" "max"
+    last_norm: None #"batch_norm", "layer_norm"
+    encoders:
+      emb_la_pos:
+        encoder_type: "laplacian_pe"
+        input_keys: ["laplacian_eigvec", "laplacian_eigval"]
+        output_keys: ["feat"]
+        hidden_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"
+      emb_rwse:
+        encoder_type: "mlp"
+        input_keys: ["rw_return_probs"]
+        output_keys: ["feat"]
+        hidden_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"
+      emb_electrostatic:
+        encoder_type: "mlp"
+        input_keys: ["electrostatic"]
+        output_keys: ["feat"]
+        hidden_dim: 32
+        num_layers: 1
+        dropout: 0.1
+        normalization: "layer_norm" #"batch_norm" or "layer_norm"
+        first_normalization: "layer_norm" #"batch_norm" or "layer_norm"
+      emb_edge_rwse:
+        encoder_type: "mlp"
+        input_keys: ["edge_rw_transition_probs"]
+        output_keys: ["edge_feat"]
+        hidden_dim: 32
+        num_layers: 1
+        dropout: 0.1
+        normalization: "layer_norm" #"batch_norm" or "layer_norm"
+      emb_edge_pes:
+        encoder_type: "cat_mlp"
+        input_keys: ["edge_rw_transition_probs", "edge_commute"]
+        output_keys: ["edge_feat"]
+        hidden_dim: 32
+        num_layers: 1
+        dropout: 0.1
+        normalization: "layer_norm" #"batch_norm" or "layer_norm"
+      gaussian_pos:
+        encoder_type: "gaussian_kernel"
+        input_keys: ["positions_3d"]
+        output_keys: ["feat", "nodepair_gaussian_bias_3d"]
+        num_heads: &num_heads 2
+        num_layers: 2
+        embed_dim: *pe_out_dim
+        use_input_keys_prefix: False
+
+  gnn:  # Set as null to avoid a post-nn network
+    out_dim: 32
+    hidden_dims: 32
+    depth: 4
+    activation: gelu
+    last_activation: none
+    dropout: 0.0
+    normalization: "layer_norm"
+    last_normalization: *normalization
+    residual_type: simple
+    pooling: [sum]
+    virtual_node: 'none'
+    layer_type: 'pyg:gps' #pyg:gine #'pyg:gps' # pyg:gated-gcn, pyg:gine,pyg:gps
+    layer_kwargs:  # Parameters for the model itself. You could define dropout_attn: 0.1
+      node_residual: false
+      mpnn_type: 'pyg:mpnnplus'
+      mpnn_kwargs:
+        in_dim: 32
+        out_dim: 32
+        in_dim_edges: 16
+        out_dim_edges: 16
+      attn_type: "full-attention" # "full-attention", "none"
+      # biased_attention: false
+      attn_kwargs:
+        num_heads: *num_heads
+      biased_attention_key: nodepair_gaussian_bias_3d
+
+
+  post_nn: null
+
+  task_heads:
+    homolumo:
+      out_dim: 1
+      hidden_dims: 256
+      depth: 2                          # Not needed if we have hidden_dims
+      activation: relu
+      last_activation: none
+      dropout: *dropout
+      normalization: *normalization
+      last_normalization: "none"
+      residual_type: none
+
+#Task-specific
+predictor:
+  metrics_on_progress_bar:
+    homolumo: ["mae", "pearsonr"]
+  loss_fun:
+    homolumo: mse_ipu
+  random_seed: *seed
+  optim_kwargs:
+    lr: 4.e-4 # warmup can be scheduled using torch_scheduler_kwargs
+    # weight_decay: 1.e-7
+  torch_scheduler_kwargs:
+    module_type: WarmUpLinearLR
+    max_num_epochs: &max_epochs 5
+    warmup_epochs: 10
+    verbose: False
+  scheduler_kwargs:
+  #  monitor: &monitor homolumo/mae/train
+  #  mode: min
+  #  frequency: 1
+  target_nan_mask: null # null: no mask, 0: 0 mask, ignore: ignore nan values from loss
+  flag_kwargs:
+    n_steps: 0 # 1
+    alpha: 0.0 # 0.01
+
+# Task-specific
+metrics:
+  homolumo:
+    - 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
+
+trainer:
+  logger:
+    save_dir: logs/PCQMv2
+    name: *name
+    project: PCQMv2_mpnn
+  #early_stopping:
+  #  monitor: *monitor
+  #  min_delta: 0
+  #  patience: 10
+  #  mode: &mode min
+  model_checkpoint:
+    dirpath: models_checkpoints/PCMQv2/
+    filename: *name
+    #monitor: *monitor
+    #mode: *mode
+    save_top_k: 1
+    every_n_epochs: 100
+  trainer:
+    precision: 32
+    max_epochs: *max_epochs
+    min_epochs: 1
+    accumulate_grad_batches: 2
+    check_val_every_n_epoch: 20
+
diff --git a/expts/configs/config_gpspp_10M_pcqm4m.yaml b/expts/configs/config_gpspp_10M_pcqm4m.yaml
new file mode 100644
index 000000000..a8d8ca5c8
--- /dev/null
+++ b/expts/configs/config_gpspp_10M_pcqm4m.yaml
@@ -0,0 +1,227 @@
+# GPS++ model with the PCQMv2 dataset on IPU.
+constants:
+  name: &name pcqm4mv2_gpspp_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
+
+datamodule:
+  module_type: "MultitaskFromSmilesDataModule"
+  # module_type: "FakeDataModule"  # Option to use generated data
+  args: # Matches that in the test_multitask_datamodule.py case.
+    task_specific_args:   # To be replaced by a new class "DatasetParams"
+      homolumo:
+        df: null
+        task_level: "graph"
+        df_path: goli/data/PCQM4Mv2/pcqm4mv2.csv #goli/data/PCQM4Mv2/pcqm4mv2.csv
+        # wget https://storage.googleapis.com/datasets-public-research/PCQM4M/cxsmiles/pcqm4mv2.csv
+        # or set path as https://storage.googleapis.com/datasets-public-research/PCQM4M/cxsmiles/pcqm4mv2.csv directly
+        smiles_col: "cxsmiles"
+        label_cols: ["homo_lumo_gap"]
+        # sample_size: 80000 # use sample_size for test
+        splits_path: goli/data/PCQM4Mv2/split_dict_v2.pt  # Download with `wget https://storage.googleapis.com/datasets-public-research/PCQM4M/cxsmiles/split_dict_v2.pt`
+        # goli/data/PCQM4Mv2/split_dict.pt
+        # goli/data/PCQM4Mv2/pcqm4m_split.csv
+        split_names: ["train", "valid", "test-dev"]
+        label_normalization:
+          method: "normal"
+          min_clipping: 0
+          max_clipping: 50
+
+    # Featurization
+    prepare_dict_or_graph: pyg:graph
+    featurization_n_jobs: 20
+    featurization_progress: True
+    featurization_backend: "loky"
+    processed_graph_data_path: "/tmp/goli_data/PCQM4Mv2/"
+    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]
+      conformer_property_list: [positions_3d] # 3D_bias
+      # 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:
+          la_pos: &pos_enc
+            pos_type: laplacian_eigvec_eigval #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
+          rw_pos: # use same name as pe_encoder
+            pos_type: rwse
+            ksteps: 16
+
+
+    # Data handling-related
+    batch_size_training: 16
+    batch_size_inference: 16
+    # cache_data_path: .
+    num_workers: 20 # -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.
+    featurization_backend: "loky"
+
+    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
+      # test-dev max nodes: 50, max_edges: 116
+      # test-challenge max nodes: 51, max_edges: 106
+
+architecture:
+  model_type: FullGraphMultiTaskNetwork
+  mup_base_path: null
+  pre_nn:   # Set as null to avoid a pre-nn network
+    out_dim: 256
+    hidden_dims: 1024
+    depth: 2
+    activation: relu
+    last_activation: none
+    dropout: &dropout 0.18
+    normalization: &normalization layer_norm
+    last_normalization: *normalization
+    residual_type: none
+
+  pre_nn_edges:   # Set as null to avoid a pre-nn network
+    out_dim: 128
+    hidden_dims: 512
+    depth: 2
+    activation: relu
+    last_activation: none
+    dropout: *dropout
+    normalization: *normalization
+    last_normalization: *normalization
+    residual_type: none
+
+  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: ["eigvecs", "eigvals"]
+        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.18
+        first_normalization: "none" #"batch_norm" or "layer_norm"
+      rw_pos:
+        encoder_type: "mlp"
+        input_keys: ["rwse"]
+        output_keys: ["feat"]
+        hidden_dim: 64
+        out_dim: 32
+        num_layers: 2
+        dropout: 0.18
+        normalization: "layer_norm" #"batch_norm" or "layer_norm"
+        first_normalization: "layer_norm" #"batch_norm" or "layer_norm"
+      gaussian_pos: # 3D_bias
+        encoder_type: "gaussian_kernel"
+        input_keys: ["positions_3d"]
+        output_keys: ["feat", "graph_gaussian_bias_3d"]
+        num_heads: 32
+        num_layers: 1 #2
+        embed_dim: 32
+        out_dim: 32 # need num of gaussian kernels 128
+        # but currently it checks pe_out_dim == pe_out_dim in encoder_manager.py, line 128
+        use_input_keys_prefix: False
+
+
+  gnn:  # Set as null to avoid a post-nn network
+    out_dim: 256
+    hidden_dims: 256
+    depth: 4
+    activation: gelu
+    last_activation: none
+    dropout: 0.1
+    normalization: "layer_norm"
+    last_normalization: *normalization
+    residual_type: simple
+    pooling: [sum]
+    virtual_node: 'none'
+    layer_type: 'pyg:gps' #pyg:gine #'pyg:gps' # pyg:gated-gcn, pyg:gine,pyg:gps
+    layer_kwargs:  # Parameters for the model itself. You could define dropout_attn: 0.1
+      node_residual: false
+      mpnn_type: 'pyg:mpnnplus'
+      mpnn_kwargs:
+        in_dim: 256
+        out_dim: 256
+        in_dim_edges: 128
+        out_dim_edges: 128
+      attn_type: "full-attention" # "full-attention", "none"
+      precision: &precision 16
+      biased_attention_key: "graph_gaussian_bias_3d" # 3D_bias
+      attn_kwargs:
+        num_heads: 32
+      droppath_rate_attn: 0.0
+      droppath_rate_ffn: 0.0
+
+
+  post_nn: null
+
+  task_heads:
+    homolumo:
+      out_dim: 1
+      hidden_dims: 256
+      depth: 2                          # Not needed if we have hidden_dims
+      activation: relu
+      last_activation: none
+      dropout: *dropout
+      normalization: *normalization
+      last_normalization: "none"
+      residual_type: none
+
+#Task-specific
+predictor:
+  metrics_on_progress_bar:
+    homolumo: ["mae", "pearsonr"]
+  loss_fun:
+    homolumo: mae_ipu
+  random_seed: *seed
+  optim_kwargs:
+    lr: 4.e-4 # warmup can be scheduled using torch_scheduler_kwargs
+    # weight_decay: 1.e-7
+    # loss_scaling: 1024
+  torch_scheduler_kwargs:
+    module_type: WarmUpLinearLR
+    max_num_epochs: &max_epochs 100
+    warmup_epochs: 10
+    verbose: False
+  scheduler_kwargs:
+  #  monitor: &monitor homolumo/mae/train
+  #  mode: min
+  #  frequency: 1
+  target_nan_mask: null # null: no mask, 0: 0 mask, ignore: ignore nan values from loss
+  flag_kwargs:
+    n_steps: 0 # 1
+    alpha: 0.0 # 0.01
+
+# Task-specific
+metrics:
+  homolumo:
+    - 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
diff --git a/expts/configs/config_mpnn_10M_b3lyp.yaml b/expts/configs/config_mpnn_10M_b3lyp.yaml
index de4591645..db7e642bd 100644
--- a/expts/configs/config_mpnn_10M_b3lyp.yaml
+++ b/expts/configs/config_mpnn_10M_b3lyp.yaml
@@ -65,6 +65,7 @@ datamodule:
     task_specific_args:   # To be replaced by a new class "DatasetParams"
       betagap:
         df: null
+        task_level: "graph"
         df_path: goli/data/b3lyp/b3lyp_mini.parquet #goli/data/b3lyp/b3lyp_mini.parquet
         # wget https://storage.googleapis.com/goli-public/datasets/b3lyp/b3lyp_mini.parquet
         # or set path as https://storage.googleapis.com/goli-public/datasets/b3lyp/b3lyp_mini.parquet directly
@@ -75,6 +76,7 @@ datamodule:
         split_test: 0.1
       alphagap:
         df: null
+        task_level: "graph"
         df_path: goli/data/b3lyp/b3lyp_mini.parquet #goli/data/b3lyp/b3lyp_mini.parquet
         # wget https://storage.googleapis.com/goli-public/datasets/b3lyp/b3lyp_mini.parquet
         # or set path as https://storage.googleapis.com/goli-public/datasets/b3lyp/b3lyp_mini.parquet directly
diff --git a/expts/configs/config_mpnn_10M_pcqm4m.yaml b/expts/configs/config_mpnn_10M_pcqm4m.yaml
index cae2c22cd..60866c776 100644
--- a/expts/configs/config_mpnn_10M_pcqm4m.yaml
+++ b/expts/configs/config_mpnn_10M_pcqm4m.yaml
@@ -55,6 +55,7 @@ datamodule:
     task_specific_args:   # To be replaced by a new class "DatasetParams"
       homolumo:
         df: null
+        task_level: "graph"
         df_path: goli/data/PCQM4M/pcqm4mv2-20k.csv
         # wget https://storage.googleapis.com/datasets-public-research/PCQM4M/cxsmiles/pcqm4mv2.csv
         # or set path as https://storage.googleapis.com/datasets-public-research/PCQM4M/cxsmiles/pcqm4mv2.csv directly
diff --git a/expts/configs/config_mpnn_pcqm4m.yaml b/expts/configs/config_mpnn_pcqm4m.yaml
index db5b1239c..046b4b738 100644
--- a/expts/configs/config_mpnn_pcqm4m.yaml
+++ b/expts/configs/config_mpnn_pcqm4m.yaml
@@ -13,6 +13,7 @@ datamodule:
     task_specific_args:   # To be replaced by a new class "DatasetParams"
       homolumo:
         df: null
+        task_level: "graph"
         df_path: goli/data/PCQM4M/pcqm4mv2-20k.csv
         # wget https://storage.googleapis.com/datasets-public-research/PCQM4M/cxsmiles/pcqm4mv2-20k.csv
         # or set path as https://storage.googleapis.com/datasets-public-research/PCQM4M/cxsmiles/pcqm4mv2-20k.csv directly
diff --git a/expts/main_run_multitask.py b/expts/main_run_multitask.py
index eee535784..3f6048754 100644
--- a/expts/main_run_multitask.py
+++ b/expts/main_run_multitask.py
@@ -1,4 +1,5 @@
 # General imports
+import argparse
 import os
 from os.path import dirname, abspath
 import yaml
@@ -21,6 +22,7 @@
     load_accelerator,
 )
 from goli.utils.safe_run import SafeRun
+from goli.utils.command_line_utils import update_config, get_anchors_and_aliases
 
 
 # WandB
@@ -30,11 +32,20 @@
 MAIN_DIR = dirname(dirname(abspath(goli.__file__)))
 
 # CONFIG_FILE = "expts/configs/config_mpnn_10M_b3lyp.yaml"
-CONFIG_FILE = "expts/configs/config_mpnn_10M_pcqm4m.yaml"
+# CONFIG_FILE = "expts/configs/config_mpnn_10M_pcqm4m.yaml"
+# CONFIG_FILE = "expts/neurips2023_configs/config_debug.yaml"
+# CONFIG_FILE = "expts/neurips2023_configs/config_large_mpnn.yaml"
+CONFIG_FILE = "expts/neurips2023_configs/config_large_gcn.yaml"
+# CONFIG_FILE = "expts/neurips2023_configs/config_large_gin.yaml"
+# CONFIG_FILE = "expts/neurips2023_configs/config_large_gine.yaml"
+# CONFIG_FILE = "expts/neurips2023_configs/config_small_gcn.yaml"
+# CONFIG_FILE = "expts/neurips2023_configs/config_small_gin.yaml"
+# CONFIG_FILE = "expts/neurips2023_configs/config_small_gine.yaml"
+
 os.chdir(MAIN_DIR)
 
 
-def main(cfg: DictConfig, run_name: str = "main", add_date_time: bool = True) -> None:
+def main(cfg: dict, run_name: str = "main", add_date_time: bool = True) -> None:
     st = timeit.default_timer()
 
     date_time_suffix = ""
@@ -42,6 +53,7 @@ def main(cfg: DictConfig, run_name: str = "main", add_date_time: bool = True) ->
         date_time_suffix = datetime.now().strftime("%d.%m.%Y_%H.%M.%S")
 
     cfg = deepcopy(cfg)
+    wandb.init(project=cfg["constants"]["name"], config=cfg)
 
     # Initialize the accelerator
     cfg, accelerator_type = load_accelerator(cfg)
@@ -93,6 +105,16 @@ def main(cfg: DictConfig, run_name: str = "main", add_date_time: bool = True) ->
 
 
 if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", help="Path to the config file", default=None)
+
+    args, unknown = parser.parse_known_args()
+    # Optionally parse the config with the command line
+    if args.config is not None:
+        CONFIG_FILE = args.config
+
     with open(os.path.join(MAIN_DIR, CONFIG_FILE), "r") as f:
         cfg = yaml.safe_load(f)
+        refs = get_anchors_and_aliases(CONFIG_FILE)
+        cfg = update_config(cfg, unknown, refs)
     main(cfg)
diff --git a/expts/neurips2023_configs/config_classifigression_l1000.yaml b/expts/neurips2023_configs/config_classifigression_l1000.yaml
new file mode 100644
index 000000000..a8bff13a3
--- /dev/null
+++ b/expts/neurips2023_configs/config_classifigression_l1000.yaml
@@ -0,0 +1,314 @@
+# Testing the mpnn only model with the PCQMv2 dataset on IPU.
+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
+ config_override:
+   datamodule:
+     batch_size_training: 64
+     batch_size_inference: 256
+   trainer:
+     trainer:
+       precision: 32
+       accumulate_grad_batches: 1
+
+datamodule:
+  module_type: "MultitaskFromSmilesDataModule"
+  # module_type: "FakeDataModule"  # Option to use generated data
+  args: # Matches that in the test_multitask_datamodule.py case.
+    task_specific_args:   # To be replaced by a new class "DatasetParams"
+      l1000_vcap:
+        df: null
+        df_path: goli/data/neurips2023/large-dataset/LINCS_L1000_VCAP_0-4.csv.gz
+        # wget https://storage.googleapis.com/goli-public/datasets/neurips_2023/Large-dataset/LINCS_L1000_VCAP_0-4.csv.gz
+        # or set path as the URL directly
+        smiles_col: "SMILES"
+        label_cols: geneID-*  # geneID-* means all columns starting with "geneID-"
+        # sample_size: 2000 # use sample_size for test
+        task_level: graph
+        splits_path: goli/data/neurips2023/small-dataset/l1000_vcap_random_splits.pt  # Download with `wget https://storage.googleapis.com/goli-public/datasets/neurips_2023/Large-dataset/l1000_vcap_random_splits.pt`
+
+      l1000_mcf7:
+        df: null
+        df_path: goli/data/neurips2023/large-dataset/LINCS_L1000_MCF7_0-4.csv.gz
+        # wget https://storage.googleapis.com/goli-public/datasets/neurips_2023/Large-dataset/LINCS_L1000_MCF7_0-4.csv.gz
+        # or set path as the URL directly
+        smiles_col: "SMILES"
+        label_cols: geneID-*  # geneID-* means all columns starting with "geneID-"
+        # sample_size: 2000 # use sample_size for test
+        task_level: graph
+        splits_path: goli/data/neurips2023/small-dataset/l1000_mcf7_random_splits.pt  # Download with `wget https://storage.googleapis.com/goli-public/datasets/neurips_2023/Large-dataset/l1000_mcf7_random_splits.pt`
+
+    # Featurization
+    prepare_dict_or_graph: pyg:graph
+    featurization_n_jobs: 1
+    featurization_progress: True
+    featurization_backend: "loky"
+    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
+
+    # cache_data_path: .
+    num_workers: 5 # -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.
+    featurization_backend: "loky"
+
+
+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:   # Set as null to avoid a pre-nn network
+    out_dim: 32
+    hidden_dims: 128
+    depth: 2
+    activation: relu
+    last_activation: none
+    dropout: *dropout
+    normalization: *normalization
+    last_normalization: *normalization
+    residual_type: none
+
+  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
+    out_dim: &gnn_dim 64
+    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:gps' #pyg:gine #'pyg:gps' # pyg:gated-gcn, pyg:gine,pyg:gps
+    layer_kwargs:  # Parameters for the model itself. You could define dropout_attn: 0.1
+      node_residual: false
+      mpnn_type: 'pyg:mpnnplus'
+      mpnn_kwargs:
+        in_dim: *gnn_dim
+        out_dim: *gnn_dim
+        in_dim_edges: 32
+        out_dim_edges: 32
+      attn_type: "none" # "full-attention", "none"
+      attn_kwargs: null
+
+
+  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:
+    l1000_vcap:
+      task_level: graph
+      out_dim: 4890  # n_columns (978) x n_brackets (5)
+      hidden_dims: 128
+      depth: 2
+      activation: none
+      last_activation: sigmoid
+      dropout: *dropout
+      normalization: *normalization
+      last_normalization: "none"
+      residual_type: none
+    l1000_mcf7:
+      task_level: graph
+      out_dim: 4890  # n_columns (978) x n_brackets (5)
+      hidden_dims: 128
+      depth: 2
+      activation: none
+      last_activation: sigmoid
+      dropout: *dropout
+      normalization: *normalization
+      last_normalization: "none"
+      residual_type: none
+
+#Task-specific
+predictor:
+  metrics_on_progress_bar:
+    l1000_vcap: []
+    l1000_mcf7: []
+  loss_fun:
+    l1000_vcap:
+      name: hybrid_ce
+      n_brackets: 5
+    l1000_mcf7:
+      name: hybrid_ce
+      n_brackets: 5
+  random_seed: *seed
+  optim_kwargs:
+    lr: 4.e-3 # warmup can be scheduled using torch_scheduler_kwargs
+    # weight_decay: 1.e-7
+  torch_scheduler_kwargs:
+    module_type: WarmUpLinearLR
+    max_num_epochs: &max_epochs 100
+    warmup_epochs: 0
+    verbose: False
+  scheduler_kwargs:
+  #  monitor: &monitor qm9/mae/train
+  #  mode: min
+  #  frequency: 1
+  target_nan_mask: ignore
+  multitask_handling: flatten
+
+# Task-specific
+metrics:
+  l1000_vcap: &classif_metrics
+    - name: auc
+      metric: auroc
+      target_nan_mask: ignore
+      multitask_handling: flatten
+      squeeze_targets: True
+      target_to_int: True
+      task: multiclass
+      num_classes: 5
+    - name: ap
+      metric: averageprecision
+      target_nan_mask: ignore
+      multitask_handling: flatten
+      squeeze_targets: True
+      target_to_int: True
+      task: multiclass
+      num_classes: 5
+    - name: accuracy
+      metric: accuracy
+      target_nan_mask: ignore
+      multitask_handling: flatten
+      squeeze_targets: True
+      target_to_int: True
+      task: multiclass
+      num_classes: 5
+      top_k: 1
+  l1000_mcf7: *classif_metrics
+
+trainer:
+  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/
+    filename: *name
+    #monitor: *monitor
+    #mode: *mode
+    save_top_k: 1
+    every_n_epochs: 100
+  trainer:
+    max_epochs: *max_epochs
+    min_epochs: 1
+    check_val_every_n_epoch: 1
diff --git a/expts/neurips2023_configs/config_debug.yaml b/expts/neurips2023_configs/config_debug.yaml
new file mode 100644
index 000000000..29bf33330
--- /dev/null
+++ b/expts/neurips2023_configs/config_debug.yaml
@@ -0,0 +1,274 @@
+# Testing the mpnn only model with the PCQMv2 dataset on IPU.
+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(8)
+#     # - 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: 16
+      batch_size_inference: 32
+    trainer:
+      trainer:
+        precision: 32
+        accumulate_grad_batches: 1
+
+datamodule:
+  module_type: "MultitaskFromSmilesDataModule"
+  # module_type: "FakeDataModule"  # Option to use generated data
+  args: # Matches that in the test_multitask_datamodule.py case.
+    task_specific_args:   # To be replaced by a new class "DatasetParams"
+
+      zinc:
+        df: null
+        df_path: goli/data/neurips2023/small-dataset/ZINC12k.csv.gz
+        # wget https://storage.googleapis.com/goli-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: 500 # use sample_size for test
+        splits_path: goli/data/neurips2023/small-dataset/ZINC12k_random_splits.pt  # Download with `wget https://storage.googleapis.com/goli-public/datasets/neurips_2023/Small-dataset/ZINC12k_random_splits.pt`
+        seed: *seed
+        task_level: graph
+        label_normalization:
+          method: "normal"
+
+    # Featurization
+    prepare_dict_or_graph: pyg:graph
+    featurization_n_jobs: 0
+    featurization_progress: True
+    featurization_backend: "loky"
+    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
+
+    # cache_data_path: .
+    num_workers: 0 # -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.
+    featurization_backend: "loky"
+
+
+architecture:
+  model_type: FullGraphMultiTaskNetwork
+  mup_base_path: null
+  pre_nn:   # Set as null to avoid a pre-nn network
+    out_dim: 32
+    hidden_dims: 32
+    depth: 2
+    activation: relu
+    last_activation: none
+    dropout: &dropout 0.18
+    normalization: &normalization layer_norm
+    last_normalization: *normalization
+    residual_type: none
+
+  pre_nn_edges:   # Set as null to avoid a pre-nn network
+    out_dim: 32
+    hidden_dims: 32
+    depth: 2
+    activation: relu
+    last_activation: none
+    dropout: *dropout
+    normalization: *normalization
+    last_normalization: *normalization
+    residual_type: none
+
+  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: 32
+        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: 32
+        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
+    out_dim: &gnn_dim 32
+    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:gps' #pyg:gine #'pyg:gps' # pyg:gated-gcn, pyg:gine,pyg:gps
+    layer_kwargs:  # Parameters for the model itself. You could define dropout_attn: 0.1
+      node_residual: false
+      mpnn_type: 'pyg:mpnnplus'
+      mpnn_kwargs:
+        in_dim: *gnn_dim
+        out_dim: *gnn_dim
+        in_dim_edges: 32
+        out_dim_edges: 32
+      attn_type: "none" # "full-attention", "none"
+      attn_kwargs: null
+
+
+  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:
+    zinc:
+      task_level: graph
+      out_dim: 3
+      hidden_dims: 32
+      depth: 2
+      activation: relu
+      last_activation: none
+      dropout: *dropout
+      normalization: *normalization
+      last_normalization: "none"
+      residual_type: none
+
+#Task-specific
+predictor:
+  metrics_on_progress_bar:
+    zinc: ["mae"]
+  loss_fun:
+    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 100
+    warmup_epochs: 10
+    verbose: False
+  scheduler_kwargs:
+  #  monitor: &monitor qm9/mae/train
+  #  mode: min
+  #  frequency: 1
+  target_nan_mask: ignore # null: no mask, 0: 0 mask, ignore-flatten, ignore-mean-per-label
+  multitask_handling: flatten # flatten, mean-per-label
+
+# Task-specific
+metrics:
+  zinc:
+    - 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
+
+trainer:
+  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/
+    filename: *name
+    #monitor: *monitor
+    #mode: *mode
+    save_top_k: 1
+    every_n_epochs: 100
+  trainer:
+    max_epochs: *max_epochs
+    min_epochs: 1
+    check_val_every_n_epoch: 20
diff --git a/expts/neurips2023_configs/config_large_gcn.yaml b/expts/neurips2023_configs/config_large_gcn.yaml
new file mode 100644
index 000000000..14d42e8f1
--- /dev/null
+++ b/expts/neurips2023_configs/config_large_gcn.yaml
@@ -0,0 +1,394 @@
+# Testing the gcn model with the PCQMv2 dataset on IPU.
+constants:
+  name: &name neurips2023_large_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: 40 # train max nodes: 20, max_edges: 54
+          max_num_edges_per_graph: 80
+        ipu_dataloader_inference_opts:
+          mode: async
+          max_num_nodes_per_graph: 40 # valid max nodes: 51, max_edges: 118
+          max_num_edges_per_graph: 80
+        # Data handling-related
+        batch_size_training: 10
+        batch_size_inference: 10
+    predictor:
+      optim_kwargs:
+        loss_scaling: 1024
+    trainer:
+      trainer:
+        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"
+  # module_type: "FakeDataModule"  # Option to use generated data
+  args: # Matches that in the test_multitask_datamodule.py case.
+    task_specific_args:   # To be replaced by a new class "DatasetParams"
+      l1000_vcap:
+        df: null
+        df_path: goli/data/neurips2023/large-dataset/LINCS_L1000_VCAP_0-4.csv.gz
+        # wget https://storage.googleapis.com/goli-public/datasets/neurips_2023/Large-dataset/LINCS_L1000_VCAP_0-4.csv.gz
+        # or set path as the URL directly
+        smiles_col: "SMILES"
+        label_cols: geneID-*  # geneID-* means all columns starting with "geneID-"
+        # sample_size: 2000 # use sample_size for test
+        task_level: graph
+        splits_path: goli/data/neurips2023/large-dataset/l1000_vcap_random_splits.pt  # Download with `wget https://storage.googleapis.com/goli-public/datasets/neurips_2023/Large-dataset/l1000_vcap_random_splits.pt`
+
+      l1000_mcf7:
+        df: null
+        df_path: goli/data/neurips2023/large-dataset/LINCS_L1000_MCF7_0-4.csv.gz
+        # wget https://storage.googleapis.com/goli-public/datasets/neurips_2023/Large-dataset/LINCS_L1000_MCF7_0-4.csv.gz
+        # or set path as the URL directly
+        smiles_col: "SMILES"
+        label_cols: geneID-*  # geneID-* means all columns starting with "geneID-"
+        # sample_size: 2000 # use sample_size for test
+        task_level: graph
+        splits_path: goli/data/neurips2023/large-dataset/l1000_mcf7_random_splits.pt  # Download with `wget https://storage.googleapis.com/goli-public/datasets/neurips_2023/Large-dataset/l1000_mcf7_random_splits.pt`
+
+      pcba_1328:
+        df: null
+        df_path: goli/data/neurips2023/large-dataset/PCBA_1328_1564k.parquet
+        # wget https://storage.googleapis.com/goli-public/datasets/neurips_2023/Large-dataset/PCBA_1328_1564k.parquet
+        # or set path as the URL directly
+        smiles_col: "SMILES"
+        label_cols: assayID-*  # assayID-* means all columns starting with "assayID-"
+        # sample_size: 2000 # use sample_size for test
+        task_level: graph
+        splits_path: goli/data/neurips2023/large-dataset/pcba_1328_random_splits.pt  # Download with `wget https://storage.googleapis.com/goli-public/datasets/neurips_2023/Large-dataset/pcba_1328_random_splits.pt`
+
+      pcqm4m_g25:
+        df: null
+        df_path: goli/data/neurips2023/large-dataset/PCQM4M_G25_N4.parquet
+        # wget https://storage.googleapis.com/goli-public/datasets/neurips_2023/Large-dataset/PCQM4M_G25_N4.parquet
+        # or set path as the URL directly
+        smiles_col: "ordered_smiles"
+        label_cols: graph_*  # graph_* means all columns starting with "graph_"
+        # sample_size: 2000 # use sample_size for test
+        task_level: graph
+        splits_path: goli/data/neurips2023/large-dataset/pcqm4m_g25_n4_random_splits.pt  # Download with `wget https://storage.googleapis.com/goli-public/datasets/neurips_2023/Large-dataset/pcqm4m_g25_n4_random_splits.pt`
+        label_normalization:
+          method: "normal"
+
+      pcqm4m_n4:
+        df: null
+        df_path: goli/data/neurips2023/large-dataset/PCQM4M_G25_N4.parquet
+        # wget https://storage.googleapis.com/goli-public/datasets/neurips_2023/Large-dataset/PCQM4M_G25_N4.parquet
+        # or set path as the URL directly
+        smiles_col: "ordered_smiles"
+        label_cols: node_* # node_* means all columns starting with "node_"
+        # sample_size: 2000 # use sample_size for test
+        task_level: node
+        splits_path: goli/data/neurips2023/large-dataset/pcqm4m_g25_n4_random_splits.pt  # Download with `wget https://storage.googleapis.com/goli-public/datasets/neurips_2023/Large-dataset/pcqm4m_g25_n4_random_splits.pt`
+        seed: *seed
+        label_normalization:
+          method: "normal"
+
+    # Featurization
+    prepare_dict_or_graph: pyg:graph
+    featurization_n_jobs: 30
+    featurization_progress: True
+    featurization_backend: "loky"
+    processed_graph_data_path: "../datacache/neurips2023-large/"
+    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
+
+    # cache_data_path: .
+    num_workers: 30 # -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.
+    featurization_backend: "loky"
+
+
+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
+
+  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 1024
+    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
+
+
+
+  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
+    node:
+      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:
+    l1000_vcap:
+      task_level: graph
+      out_dim: 4890
+      hidden_dims: 128
+      depth: 2
+      activation: none
+      last_activation: sigmoid
+      dropout: *dropout
+      normalization: *normalization
+      last_normalization: "none"
+      residual_type: none
+    l1000_mcf7:
+      task_level: graph
+      out_dim: 4890
+      hidden_dims: 128
+      depth: 2
+      activation: none
+      last_activation: sigmoid
+      dropout: *dropout
+      normalization: *normalization
+      last_normalization: "none"
+      residual_type: none
+    pcba_1328:
+      task_level: graph
+      out_dim: 1328
+      hidden_dims: 64
+      depth: 2
+      activation: relu
+      last_activation: sigmoid
+      dropout: *dropout
+      normalization: *normalization
+      last_normalization: "none"
+      residual_type: none
+    pcqm4m_g25:
+      task_level: graph
+      out_dim: 25
+      hidden_dims: 32
+      depth: 2
+      activation: relu
+      last_activation: none
+      dropout: *dropout
+      normalization: *normalization
+      last_normalization: "none"
+      residual_type: none
+    pcqm4m_n4:
+      task_level: node
+      out_dim: 4
+      hidden_dims: 32
+      depth: 2
+      activation: relu
+      last_activation: none
+      dropout: *dropout
+      normalization: *normalization
+      last_normalization: "none"
+      residual_type: none
+
+#Task-specific
+predictor:
+  metrics_on_progress_bar:
+    l1000_vcap: []
+    l1000_mcf7: []
+    pcba_1328: []
+    pcqm4m_g25: []
+    pcqm4m_n4: []
+  loss_fun:
+    l1000_vcap: hybrid_ce_ipu
+    l1000_mcf7: hybrid_ce_ipu
+    pcba_1328: bce_ipu
+    pcqm4m_g25: mae_ipu
+    pcqm4m_n4: 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 100
+    warmup_epochs: 10
+    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:
+  l1000_vcap: &classif_metrics
+    - name: auroc_flat
+      metric: auroc_ipu
+      num_classes: 5
+      task: multiclass
+      multitask_handling: flatten
+      threshold_kwargs: null
+    - name: avpr_flat
+      metric: average_precision_ipu
+      num_classes: 5
+      task: multiclass
+      target_to_int: True
+      target_nan_mask: -1000
+      ignore_index: -1000
+      multitask_handling: flatten
+      threshold_kwargs: null
+  l1000_mcf7: *classif_metrics
+  pcba_1328:
+    - name: auroc_flat
+      metric: auroc_ipu
+      task: binary
+      multitask_handling: flatten
+      threshold_kwargs: null
+    - name: avpr_flat
+      metric: average_precision_ipu
+      task: binary
+      multitask_handling: flatten
+      threshold_kwargs: null
+  pcqm4m_g25: &pcqm_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
+  pcqm4m_n4: *pcqm_metrics
+
+trainer:
+  logger:
+    save_dir: logs/neurips2023-large/
+    name: *name
+    project: *name
+  #early_stopping:
+  #  monitor: *monitor
+  #  min_delta: 0
+  #  patience: 10
+  #  mode: &mode min
+  model_checkpoint:
+    dirpath: models_checkpoints/neurips2023-large/
+    filename: *name
+    #monitor: *monitor
+    #mode: *mode
+    save_top_k: 1
+    every_n_epochs: 100
+  trainer:
+    max_epochs: *max_epochs
+    min_epochs: 1
+    check_val_every_n_epoch: 20
diff --git a/expts/neurips2023_configs/config_large_gin.yaml b/expts/neurips2023_configs/config_large_gin.yaml
new file mode 100644
index 000000000..f5dd74028
--- /dev/null
+++ b/expts/neurips2023_configs/config_large_gin.yaml
@@ -0,0 +1,394 @@
+# Testing the gin model with the PCQMv2 dataset on IPU.
+constants:
+  name: &name neurips2023_large_data_gin
+  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: 40 # train max nodes: 20, max_edges: 54
+          max_num_edges_per_graph: 80
+        ipu_dataloader_inference_opts:
+          mode: async
+          max_num_nodes_per_graph: 40 # valid max nodes: 51, max_edges: 118
+          max_num_edges_per_graph: 80
+        # Data handling-related
+        batch_size_training: 10
+        batch_size_inference: 10
+    predictor:
+      optim_kwargs:
+        loss_scaling: 1024
+    trainer:
+      trainer:
+        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"
+  # module_type: "FakeDataModule"  # Option to use generated data
+  args: # Matches that in the test_multitask_datamodule.py case.
+    task_specific_args:   # To be replaced by a new class "DatasetParams"
+      l1000_vcap:
+        df: null
+        df_path: goli/data/neurips2023/large-dataset/LINCS_L1000_VCAP_0-4.csv.gz
+        # wget https://storage.googleapis.com/goli-public/datasets/neurips_2023/Large-dataset/LINCS_L1000_VCAP_0-4.csv.gz
+        # or set path as the URL directly
+        smiles_col: "SMILES"
+        label_cols: geneID-*  # geneID-* means all columns starting with "geneID-"
+        # sample_size: 2000 # use sample_size for test
+        task_level: graph
+        splits_path: goli/data/neurips2023/large-dataset/l1000_vcap_random_splits.pt  # Download with `wget https://storage.googleapis.com/goli-public/datasets/neurips_2023/Large-dataset/l1000_vcap_random_splits.pt`
+
+      l1000_mcf7:
+        df: null
+        df_path: goli/data/neurips2023/large-dataset/LINCS_L1000_MCF7_0-4.csv.gz
+        # wget https://storage.googleapis.com/goli-public/datasets/neurips_2023/Large-dataset/LINCS_L1000_MCF7_0-4.csv.gz
+        # or set path as the URL directly
+        smiles_col: "SMILES"
+        label_cols: geneID-*  # geneID-* means all columns starting with "geneID-"
+        # sample_size: 2000 # use sample_size for test
+        task_level: graph
+        splits_path: goli/data/neurips2023/large-dataset/l1000_mcf7_random_splits.pt  # Download with `wget https://storage.googleapis.com/goli-public/datasets/neurips_2023/Large-dataset/l1000_mcf7_random_splits.pt`
+
+      pcba_1328:
+        df: null
+        df_path: goli/data/neurips2023/large-dataset/PCBA_1328_1564k.parquet
+        # wget https://storage.googleapis.com/goli-public/datasets/neurips_2023/Large-dataset/PCBA_1328_1564k.parquet
+        # or set path as the URL directly
+        smiles_col: "SMILES"
+        label_cols: assayID-*  # assayID-* means all columns starting with "assayID-"
+        # sample_size: 2000 # use sample_size for test
+        task_level: graph
+        splits_path: goli/data/neurips2023/large-dataset/pcba_1328_random_splits.pt  # Download with `wget https://storage.googleapis.com/goli-public/datasets/neurips_2023/Large-dataset/pcba_1328_random_splits.pt`
+
+      pcqm4m_g25:
+        df: null
+        df_path: goli/data/neurips2023/large-dataset/PCQM4M_G25_N4.parquet
+        # wget https://storage.googleapis.com/goli-public/datasets/neurips_2023/Large-dataset/PCQM4M_G25_N4.parquet
+        # or set path as the URL directly
+        smiles_col: "ordered_smiles"
+        label_cols: graph_*  # graph_* means all columns starting with "graph_"
+        # sample_size: 2000 # use sample_size for test
+        task_level: graph
+        splits_path: goli/data/neurips2023/large-dataset/pcqm4m_g25_n4_random_splits.pt  # Download with `wget https://storage.googleapis.com/goli-public/datasets/neurips_2023/Large-dataset/pcqm4m_g25_n4_random_splits.pt`
+        label_normalization:
+          method: "normal"
+
+      pcqm4m_n4:
+        df: null
+        df_path: goli/data/neurips2023/large-dataset/PCQM4M_G25_N4.parquet
+        # wget https://storage.googleapis.com/goli-public/datasets/neurips_2023/Large-dataset/PCQM4M_G25_N4.parquet
+        # or set path as the URL directly
+        smiles_col: "ordered_smiles"
+        label_cols: node_* # node_* means all columns starting with "node_"
+        # sample_size: 2000 # use sample_size for test
+        task_level: node
+        splits_path: goli/data/neurips2023/large-dataset/pcqm4m_g25_n4_random_splits.pt  # Download with `wget https://storage.googleapis.com/goli-public/datasets/neurips_2023/Large-dataset/pcqm4m_g25_n4_random_splits.pt`
+        seed: *seed
+        label_normalization:
+          method: "normal"
+
+    # Featurization
+    prepare_dict_or_graph: pyg:graph
+    featurization_n_jobs: 30
+    featurization_progress: True
+    featurization_backend: "loky"
+    processed_graph_data_path: "../datacache/neurips2023-large/"
+    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
+
+    # cache_data_path: .
+    num_workers: 30 # -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.
+    featurization_backend: "loky"
+
+
+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
+
+  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 832
+    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:gin' #pyg:gine #'pyg:gps' # pyg:gated-gcn, pyg:gine,pyg:gps
+
+
+
+  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
+    node:
+      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:
+    l1000_vcap:
+      task_level: graph
+      out_dim: 4890
+      hidden_dims: 128
+      depth: 2
+      activation: none
+      last_activation: sigmoid
+      dropout: *dropout
+      normalization: *normalization
+      last_normalization: "none"
+      residual_type: none
+    l1000_mcf7:
+      task_level: graph
+      out_dim: 4890
+      hidden_dims: 128
+      depth: 2
+      activation: none
+      last_activation: sigmoid
+      dropout: *dropout
+      normalization: *normalization
+      last_normalization: "none"
+      residual_type: none
+    pcba_1328:
+      task_level: graph
+      out_dim: 1328
+      hidden_dims: 64
+      depth: 2
+      activation: relu
+      last_activation: sigmoid
+      dropout: *dropout
+      normalization: *normalization
+      last_normalization: "none"
+      residual_type: none
+    pcqm4m_g25:
+      task_level: graph
+      out_dim: 25
+      hidden_dims: 32
+      depth: 2
+      activation: relu
+      last_activation: none
+      dropout: *dropout
+      normalization: *normalization
+      last_normalization: "none"
+      residual_type: none
+    pcqm4m_n4:
+      task_level: node
+      out_dim: 4
+      hidden_dims: 32
+      depth: 2
+      activation: relu
+      last_activation: none
+      dropout: *dropout
+      normalization: *normalization
+      last_normalization: "none"
+      residual_type: none
+
+#Task-specific
+predictor:
+  metrics_on_progress_bar:
+    l1000_vcap: []
+    l1000_mcf7: []
+    pcba_1328: []
+    pcqm4m_g25: []
+    pcqm4m_n4: []
+  loss_fun:
+    l1000_vcap: hybrid_ce_ipu
+    l1000_mcf7: hybrid_ce_ipu
+    pcba_1328: bce_ipu
+    pcqm4m_g25: mae_ipu
+    pcqm4m_n4: 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 100
+    warmup_epochs: 10
+    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:
+  l1000_vcap: &classif_metrics
+    - name: auroc_flat
+      metric: auroc_ipu
+      num_classes: 5
+      task: multiclass
+      multitask_handling: flatten
+      threshold_kwargs: null
+    - name: avpr_flat
+      metric: average_precision_ipu
+      num_classes: 5
+      task: multiclass
+      target_to_int: True
+      target_nan_mask: -1000
+      ignore_index: -1000
+      multitask_handling: flatten
+      threshold_kwargs: null
+  l1000_mcf7: *classif_metrics
+  pcba_1328:
+    - name: auroc_flat
+      metric: auroc_ipu
+      task: binary
+      multitask_handling: flatten
+      threshold_kwargs: null
+    - name: avpr_flat
+      metric: average_precision_ipu
+      task: binary
+      multitask_handling: flatten
+      threshold_kwargs: null
+  pcqm4m_g25: &pcqm_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
+  pcqm4m_n4: *pcqm_metrics
+
+trainer:
+  logger:
+    save_dir: logs/neurips2023-large/
+    name: *name
+    project: *name
+  #early_stopping:
+  #  monitor: *monitor
+  #  min_delta: 0
+  #  patience: 10
+  #  mode: &mode min
+  model_checkpoint:
+    dirpath: models_checkpoints/neurips2023-large/
+    filename: *name
+    #monitor: *monitor
+    #mode: *mode
+    save_top_k: 1
+    every_n_epochs: 100
+  trainer:
+    max_epochs: *max_epochs
+    min_epochs: 1
+    check_val_every_n_epoch: 20
diff --git a/expts/neurips2023_configs/config_large_gine.yaml b/expts/neurips2023_configs/config_large_gine.yaml
new file mode 100644
index 000000000..d8f123a22
--- /dev/null
+++ b/expts/neurips2023_configs/config_large_gine.yaml
@@ -0,0 +1,401 @@
+# Testing the gine model with the PCQMv2 dataset on IPU.
+constants:
+  name: &name neurips2023_large_data_gine
+  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: 40 # train max nodes: 20, max_edges: 54
+          max_num_edges_per_graph: 80
+        ipu_dataloader_inference_opts:
+          mode: async
+          max_num_nodes_per_graph: 40 # valid max nodes: 51, max_edges: 118
+          max_num_edges_per_graph: 80
+        # Data handling-related
+        batch_size_training: 10
+        batch_size_inference: 10
+    predictor:
+      optim_kwargs:
+        loss_scaling: 1024
+    trainer:
+      trainer:
+        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"
+  # module_type: "FakeDataModule"  # Option to use generated data
+  args: # Matches that in the test_multitask_datamodule.py case.
+    task_specific_args:   # To be replaced by a new class "DatasetParams"
+      l1000_vcap:
+        df: null
+        df_path: goli/data/neurips2023/large-dataset/LINCS_L1000_VCAP_0-4.csv.gz
+        # wget https://storage.googleapis.com/goli-public/datasets/neurips_2023/Large-dataset/LINCS_L1000_VCAP_0-4.csv.gz
+        # or set path as the URL directly
+        smiles_col: "SMILES"
+        label_cols: geneID-*  # geneID-* means all columns starting with "geneID-"
+        # sample_size: 2000 # use sample_size for test
+        task_level: graph
+        splits_path: goli/data/neurips2023/large-dataset/l1000_vcap_random_splits.pt  # Download with `wget https://storage.googleapis.com/goli-public/datasets/neurips_2023/Large-dataset/l1000_vcap_random_splits.pt`
+
+      l1000_mcf7:
+        df: null
+        df_path: goli/data/neurips2023/large-dataset/LINCS_L1000_MCF7_0-4.csv.gz
+        # wget https://storage.googleapis.com/goli-public/datasets/neurips_2023/Large-dataset/LINCS_L1000_MCF7_0-4.csv.gz
+        # or set path as the URL directly
+        smiles_col: "SMILES"
+        label_cols: geneID-*  # geneID-* means all columns starting with "geneID-"
+        # sample_size: 2000 # use sample_size for test
+        task_level: graph
+        splits_path: goli/data/neurips2023/large-dataset/l1000_mcf7_random_splits.pt  # Download with `wget https://storage.googleapis.com/goli-public/datasets/neurips_2023/Large-dataset/l1000_mcf7_random_splits.pt`
+
+      pcba_1328:
+        df: null
+        df_path: goli/data/neurips2023/large-dataset/PCBA_1328_1564k.parquet
+        # wget https://storage.googleapis.com/goli-public/datasets/neurips_2023/Large-dataset/PCBA_1328_1564k.parquet
+        # or set path as the URL directly
+        smiles_col: "SMILES"
+        label_cols: assayID-*  # assayID-* means all columns starting with "assayID-"
+        # sample_size: 2000 # use sample_size for test
+        task_level: graph
+        splits_path: goli/data/neurips2023/large-dataset/pcba_1328_random_splits.pt  # Download with `wget https://storage.googleapis.com/goli-public/datasets/neurips_2023/Large-dataset/pcba_1328_random_splits.pt`
+
+      pcqm4m_g25:
+        df: null
+        df_path: goli/data/neurips2023/large-dataset/PCQM4M_G25_N4.parquet
+        # wget https://storage.googleapis.com/goli-public/datasets/neurips_2023/Large-dataset/PCQM4M_G25_N4.parquet
+        # or set path as the URL directly
+        smiles_col: "ordered_smiles"
+        label_cols: graph_*  # graph_* means all columns starting with "graph_"
+        # sample_size: 2000 # use sample_size for test
+        task_level: graph
+        splits_path: goli/data/neurips2023/large-dataset/pcqm4m_g25_n4_random_splits.pt  # Download with `wget https://storage.googleapis.com/goli-public/datasets/neurips_2023/Large-dataset/pcqm4m_g25_n4_random_splits.pt`
+        label_normalization:
+          method: "normal"
+
+      pcqm4m_n4:
+        df: null
+        df_path: goli/data/neurips2023/large-dataset/PCQM4M_G25_N4.parquet
+        # wget https://storage.googleapis.com/goli-public/datasets/neurips_2023/Large-dataset/PCQM4M_G25_N4.parquet
+        # or set path as the URL directly
+        smiles_col: "ordered_smiles"
+        label_cols: node_* # node_* means all columns starting with "node_"
+        # sample_size: 2000 # use sample_size for test
+        task_level: node
+        splits_path: goli/data/neurips2023/large-dataset/pcqm4m_g25_n4_random_splits.pt  # Download with `wget https://storage.googleapis.com/goli-public/datasets/neurips_2023/Large-dataset/pcqm4m_g25_n4_random_splits.pt`
+        seed: *seed
+        label_normalization:
+          method: "normal"
+
+    # Featurization
+    prepare_dict_or_graph: pyg:graph
+    featurization_n_jobs: 30
+    featurization_progress: True
+    featurization_backend: "loky"
+    processed_graph_data_path: "../datacache/neurips2023-large/"
+    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
+
+    # cache_data_path: .
+    num_workers: 30 # -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.
+    featurization_backend: "loky"
+
+
+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:   # Set as null to avoid a pre-nn network
+    out_dim: 32
+    hidden_dims: 128
+    depth: 2
+    activation: relu
+    last_activation: none
+    dropout: *dropout
+    normalization: *normalization
+    last_normalization: *normalization
+    residual_type: none
+
+  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
+    out_dim: &gnn_dim 832
+    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:gine' #pyg:gine #'pyg:gps' # pyg:gated-gcn, pyg:gine,pyg:gps
+
+
+  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
+    node:
+      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:
+    l1000_vcap:
+      task_level: graph
+      out_dim: 4890
+      hidden_dims: 128
+      depth: 2
+      activation: none
+      last_activation: sigmoid
+      dropout: *dropout
+      normalization: *normalization
+      last_normalization: "none"
+      residual_type: none
+    l1000_mcf7:
+      task_level: graph
+      out_dim: 4890
+      hidden_dims: 128
+      depth: 2
+      activation: none
+      last_activation: sigmoid
+      dropout: *dropout
+      normalization: *normalization
+      last_normalization: "none"
+      residual_type: none
+    pcba_1328:
+      task_level: graph
+      out_dim: 1328
+      hidden_dims: 64
+      depth: 2
+      activation: relu
+      last_activation: sigmoid
+      dropout: *dropout
+      normalization: *normalization
+      last_normalization: "none"
+      residual_type: none
+    pcqm4m_g25:
+      task_level: graph
+      out_dim: 25
+      hidden_dims: 32
+      depth: 2
+      activation: relu
+      last_activation: none
+      dropout: *dropout
+      normalization: *normalization
+      last_normalization: "none"
+      residual_type: none
+    pcqm4m_n4:
+      task_level: node
+      out_dim: 4
+      hidden_dims: 32
+      depth: 2
+      activation: relu
+      last_activation: none
+      dropout: *dropout
+      normalization: *normalization
+      last_normalization: "none"
+      residual_type: none
+
+#Task-specific
+predictor:
+  metrics_on_progress_bar:
+    l1000_vcap: []
+    l1000_mcf7: []
+    pcba_1328: []
+    pcqm4m_g25: []
+    pcqm4m_n4: []
+  loss_fun:
+    l1000_vcap: hybrid_ce_ipu
+    l1000_mcf7: hybrid_ce_ipu
+    pcba_1328: bce_ipu
+    pcqm4m_g25: mae_ipu
+    pcqm4m_n4: 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 100
+    warmup_epochs: 10
+    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:
+  l1000_vcap: &classif_metrics
+    - name: auroc_flat
+      metric: auroc_ipu
+      num_classes: 5
+      task: multiclass
+      multitask_handling: flatten
+      threshold_kwargs: null
+    - name: avpr_flat
+      metric: average_precision_ipu
+      num_classes: 5
+      task: multiclass
+      target_to_int: True
+      target_nan_mask: -1000
+      ignore_index: -1000
+      multitask_handling: flatten
+      threshold_kwargs: null
+  l1000_mcf7: *classif_metrics
+  pcba_1328:
+    - name: auroc_flat
+      metric: auroc_ipu
+      task: binary
+      multitask_handling: flatten
+      threshold_kwargs: null
+    - name: avpr_flat
+      metric: average_precision_ipu
+      task: binary
+      multitask_handling: flatten
+      threshold_kwargs: null
+  pcqm4m_g25: &pcqm_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
+  pcqm4m_n4: *pcqm_metrics
+
+trainer:
+  logger:
+    save_dir: logs/neurips2023-large/
+    name: *name
+    project: *name
+  #early_stopping:
+  #  monitor: *monitor
+  #  min_delta: 0
+  #  patience: 10
+  #  mode: &mode min
+  model_checkpoint:
+    dirpath: models_checkpoints/neurips2023-large/
+    filename: *name
+    #monitor: *monitor
+    #mode: *mode
+    save_top_k: 1
+    every_n_epochs: 100
+  trainer:
+    max_epochs: *max_epochs
+    min_epochs: 1
+    check_val_every_n_epoch: 20
diff --git a/expts/neurips2023_configs/config_large_mpnn.yaml b/expts/neurips2023_configs/config_large_mpnn.yaml
new file mode 100644
index 000000000..0845e2eb4
--- /dev/null
+++ b/expts/neurips2023_configs/config_large_mpnn.yaml
@@ -0,0 +1,411 @@
+# Testing the mpnn only model with the PCQMv2 dataset on IPU.
+constants:
+  name: &name neurips2023_large_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: 40 # train max nodes: 20, max_edges: 54
+          max_num_edges_per_graph: 80
+        ipu_dataloader_inference_opts:
+          mode: async
+          max_num_nodes_per_graph: 40 # valid max nodes: 51, max_edges: 118
+          max_num_edges_per_graph: 80
+        # Data handling-related
+        batch_size_training: 10
+        batch_size_inference: 10
+    predictor:
+      optim_kwargs:
+        loss_scaling: 1024
+    trainer:
+      trainer:
+        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"
+  # module_type: "FakeDataModule"  # Option to use generated data
+  args: # Matches that in the test_multitask_datamodule.py case.
+    task_specific_args:   # To be replaced by a new class "DatasetParams"
+      l1000_vcap:
+        df: null
+        df_path: goli/data/neurips2023/large-dataset/LINCS_L1000_VCAP_0-4.csv.gz
+        # wget https://storage.googleapis.com/goli-public/datasets/neurips_2023/Large-dataset/LINCS_L1000_VCAP_0-4.csv.gz
+        # or set path as the URL directly
+        smiles_col: "SMILES"
+        label_cols: geneID-*  # geneID-* means all columns starting with "geneID-"
+        sample_size: 20000 # use sample_size for test
+        task_level: graph
+        splits_path: goli/data/neurips2023/large-dataset/l1000_vcap_random_splits.pt  # Download with `wget https://storage.googleapis.com/goli-public/datasets/neurips_2023/Large-dataset/l1000_vcap_random_splits.pt`
+
+      l1000_mcf7:
+        df: null
+        df_path: goli/data/neurips2023/large-dataset/LINCS_L1000_MCF7_0-4.csv.gz
+        # wget https://storage.googleapis.com/goli-public/datasets/neurips_2023/Large-dataset/LINCS_L1000_MCF7_0-4.csv.gz
+        # or set path as the URL directly
+        smiles_col: "SMILES"
+        label_cols: geneID-*  # geneID-* means all columns starting with "geneID-"
+        sample_size: 20000 # use sample_size for test
+        task_level: graph
+        splits_path: goli/data/neurips2023/large-dataset/l1000_mcf7_random_splits.pt  # Download with `wget https://storage.googleapis.com/goli-public/datasets/neurips_2023/Large-dataset/l1000_mcf7_random_splits.pt`
+
+      pcba_1328:
+        df: null
+        df_path: goli/data/neurips2023/large-dataset/PCBA_1328_1564k.parquet
+        # wget https://storage.googleapis.com/goli-public/datasets/neurips_2023/Large-dataset/PCBA_1328_1564k.parquet
+        # or set path as the URL directly
+        smiles_col: "SMILES"
+        label_cols: assayID-*  # assayID-* means all columns starting with "assayID-"
+        sample_size: 20000 # use sample_size for test
+        task_level: graph
+        splits_path: goli/data/neurips2023/large-dataset/pcba_1328_random_splits.pt  # Download with `wget https://storage.googleapis.com/goli-public/datasets/neurips_2023/Large-dataset/pcba_1328_random_splits.pt`
+
+      pcqm4m_g25:
+        df: null
+        df_path: goli/data/neurips2023/large-dataset/PCQM4M_G25_N4.parquet
+        # wget https://storage.googleapis.com/goli-public/datasets/neurips_2023/Large-dataset/PCQM4M_G25_N4.parquet
+        # or set path as the URL directly
+        smiles_col: "ordered_smiles"
+        label_cols: graph_*  # graph_* means all columns starting with "graph_"
+        sample_size: 20000 # use sample_size for test
+        task_level: graph
+        splits_path: goli/data/neurips2023/large-dataset/pcqm4m_g25_n4_random_splits.pt  # Download with `wget https://storage.googleapis.com/goli-public/datasets/neurips_2023/Large-dataset/pcqm4m_g25_n4_random_splits.pt`
+        label_normalization:
+          method: "normal"
+
+      pcqm4m_n4:
+        df: null
+        df_path: goli/data/neurips2023/large-dataset/PCQM4M_G25_N4.parquet
+        # wget https://storage.googleapis.com/goli-public/datasets/neurips_2023/Large-dataset/PCQM4M_G25_N4.parquet
+        # or set path as the URL directly
+        smiles_col: "ordered_smiles"
+        label_cols: node_* # node_* means all columns starting with "node_"
+        sample_size: 20000 # use sample_size for test
+        task_level: node
+        splits_path: goli/data/neurips2023/large-dataset/pcqm4m_g25_n4_random_splits.pt  # Download with `wget https://storage.googleapis.com/goli-public/datasets/neurips_2023/Large-dataset/pcqm4m_g25_n4_random_splits.pt`
+        seed: *seed
+        label_normalization:
+          method: "normal"
+
+    # Featurization
+    prepare_dict_or_graph: pyg:graph
+    featurization_n_jobs: 30
+    featurization_progress: True
+    featurization_backend: "loky"
+    processed_graph_data_path: "../datacache/neurips2023-large/"
+    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
+
+    # cache_data_path: .
+    num_workers: 30 # -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.
+    featurization_backend: "loky"
+
+
+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:   # Set as null to avoid a pre-nn network
+    out_dim: 32
+    hidden_dims: 128
+    depth: 2
+    activation: relu
+    last_activation: none
+    dropout: *dropout
+    normalization: *normalization
+    last_normalization: *normalization
+    residual_type: none
+
+  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
+    out_dim: &gnn_dim 64
+    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:gps' #pyg:gine #'pyg:gps' # pyg:gated-gcn, pyg:gine,pyg:gps
+    layer_kwargs:  # Parameters for the model itself. You could define dropout_attn: 0.1
+      node_residual: false
+      mpnn_type: 'pyg:mpnnplus'
+      mpnn_kwargs:
+        in_dim: *gnn_dim
+        out_dim: *gnn_dim
+        in_dim_edges: 32
+        out_dim_edges: 32
+      attn_type: "none" # "full-attention", "none"
+      attn_kwargs: null
+
+
+  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
+    node:
+      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:
+    l1000_vcap:
+      task_level: graph
+      out_dim: 4890
+      hidden_dims: 128
+      depth: 2
+      activation: none
+      last_activation: sigmoid
+      dropout: *dropout
+      normalization: *normalization
+      last_normalization: "none"
+      residual_type: none
+    l1000_mcf7:
+      task_level: graph
+      out_dim: 4890
+      hidden_dims: 128
+      depth: 2
+      activation: none
+      last_activation: sigmoid
+      dropout: *dropout
+      normalization: *normalization
+      last_normalization: "none"
+      residual_type: none
+    pcba_1328:
+      task_level: graph
+      out_dim: 1328
+      hidden_dims: 64
+      depth: 2
+      activation: relu
+      last_activation: sigmoid
+      dropout: *dropout
+      normalization: *normalization
+      last_normalization: "none"
+      residual_type: none
+    pcqm4m_g25:
+      task_level: graph
+      out_dim: 25
+      hidden_dims: 32
+      depth: 2
+      activation: relu
+      last_activation: none
+      dropout: *dropout
+      normalization: *normalization
+      last_normalization: "none"
+      residual_type: none
+    pcqm4m_n4:
+      task_level: node
+      out_dim: 4
+      hidden_dims: 32
+      depth: 2
+      activation: relu
+      last_activation: none
+      dropout: *dropout
+      normalization: *normalization
+      last_normalization: "none"
+      residual_type: none
+
+#Task-specific
+predictor:
+  metrics_on_progress_bar:
+    l1000_vcap: []
+    l1000_mcf7: []
+    pcba_1328: []
+    pcqm4m_g25: []
+    pcqm4m_n4: []
+  loss_fun:
+    l1000_vcap: hybrid_ce_ipu
+    l1000_mcf7: hybrid_ce_ipu
+    pcba_1328: bce_ipu
+    pcqm4m_g25: mae_ipu
+    pcqm4m_n4: 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 100
+    warmup_epochs: 10
+    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:
+  l1000_vcap: &classif_metrics
+    - name: auroc_flat
+      metric: auroc_ipu
+      num_classes: 5
+      task: multiclass
+      multitask_handling: flatten
+      threshold_kwargs: null
+    - name: avpr_flat
+      metric: average_precision_ipu
+      num_classes: 5
+      task: multiclass
+      target_to_int: True
+      target_nan_mask: -1000
+      ignore_index: -1000
+      multitask_handling: flatten
+      threshold_kwargs: null
+  l1000_mcf7: *classif_metrics
+  pcba_1328:
+    - name: auroc_flat
+      metric: auroc_ipu
+      task: binary
+      multitask_handling: flatten
+      threshold_kwargs: null
+    - name: avpr_flat
+      metric: average_precision_ipu
+      task: binary
+      multitask_handling: flatten
+      threshold_kwargs: null
+  pcqm4m_g25: &pcqm_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
+  pcqm4m_n4: *pcqm_metrics
+
+trainer:
+  logger:
+    save_dir: logs/neurips2023-large/
+    name: *name
+    project: *name
+  #early_stopping:
+  #  monitor: *monitor
+  #  min_delta: 0
+  #  patience: 10
+  #  mode: &mode min
+  model_checkpoint:
+    dirpath: models_checkpoints/neurips2023-large/
+    filename: *name
+    #monitor: *monitor
+    #mode: *mode
+    save_top_k: 1
+    every_n_epochs: 100
+  trainer:
+    max_epochs: *max_epochs
+    min_epochs: 1
+    check_val_every_n_epoch: 20
diff --git a/expts/neurips2023_configs/config_luis_jama.yaml b/expts/neurips2023_configs/config_luis_jama.yaml
new file mode 100644
index 000000000..3a17f74c2
--- /dev/null
+++ b/expts/neurips2023_configs/config_luis_jama.yaml
@@ -0,0 +1,314 @@
+# Testing the mpnn only model with the PCQMv2 dataset on IPU.
+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
+  config_override:
+    datamodule:
+      batch_size_training: 64
+      batch_size_inference: 256
+    trainer:
+      trainer:
+        precision: 32
+        accumulate_grad_batches: 1
+
+datamodule:
+  module_type: "MultitaskFromSmilesDataModule"
+  # module_type: "FakeDataModule"  # Option to use generated data
+  args: # Matches that in the test_multitask_datamodule.py case.
+    task_specific_args:   # To be replaced by a new class "DatasetParams"
+
+      pcqm20k_g13:
+        df: null
+        df_path: goli/data/neurips2023/dummy-dataset/PCQM20k_G13_N4.parquet
+        # wget https://storage.googleapis.com/goli-public/datasets/neurips_2023/Dummy-dataset/PCQM20k_G13_N4.parquet
+        # or set path as the URL directly
+        smiles_col: "ordered_smiles"
+        label_cols: graph_*  # graph_* means all columns starting with "graph_"
+        sample_size: 2000 # use sample_size for test
+        split_val: 0.1
+        split_test: 0.1
+        task_level: graph
+        label_normalization:
+          method: "normal"
+
+      pcqm20k_n4:
+        df: null
+        df_path: goli/data/neurips2023/dummy-dataset/PCQM20k_G13_N4.parquet
+        # wget https://storage.googleapis.com/goli-public/datasets/neurips_2023/Dummy-dataset/PCQM20k_G13_N4.parquet
+        # or set path as the URL directly
+        smiles_col: "ordered_smiles"
+        label_cols: node_* # node_* means all columns starting with "node_"
+        sample_size: 2000 # use sample_size for test
+        split_val: 0.1
+        split_test: 0.1
+        seed: *seed
+        task_level: node
+        label_normalization:
+          method: "normal"
+
+    # Featurization
+    prepare_dict_or_graph: pyg:graph
+    featurization_n_jobs: 30
+    featurization_progress: True
+    featurization_backend: "loky"
+    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
+
+    # cache_data_path: .
+    num_workers: 4 # -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.
+    featurization_backend: "loky"
+
+
+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:   # Set as null to avoid a pre-nn network
+    out_dim: 32
+    hidden_dims: 128
+    depth: 2
+    activation: relu
+    last_activation: none
+    dropout: *dropout
+    normalization: *normalization
+    last_normalization: *normalization
+    residual_type: none
+
+  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
+    out_dim: &gnn_dim 64
+    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:gps' #pyg:gine #'pyg:gps' # pyg:gated-gcn, pyg:gine,pyg:gps
+    layer_kwargs:  # Parameters for the model itself. You could define dropout_attn: 0.1
+      node_residual: false
+      mpnn_type: 'pyg:mpnnplus'
+      mpnn_kwargs:
+        in_dim: *gnn_dim
+        out_dim: *gnn_dim
+        in_dim_edges: 32
+        out_dim_edges: 32
+      attn_type: "none" # "full-attention", "none"
+      attn_kwargs: null
+
+
+  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
+    node:
+      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:
+    pcqm20k_g13:
+      task_level: graph
+      out_dim: 13
+      hidden_dims: 32
+      depth: 2
+      activation: relu
+      last_activation: none
+      dropout: *dropout
+      normalization: *normalization
+      last_normalization: "none"
+      residual_type: none
+    pcqm20k_n4:
+      task_level: node
+      out_dim: 4
+      hidden_dims: 32
+      depth: 2
+      activation: relu
+      last_activation: none
+      dropout: *dropout
+      normalization: *normalization
+      last_normalization: "none"
+      residual_type: none
+
+#Task-specific
+predictor:
+  metrics_on_progress_bar:
+    pcqm20k_g13: []
+    pcqm20k_n4: []
+  loss_fun:
+    pcqm20k_g13: mae_ipu
+    pcqm20k_n4: 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 100
+    warmup_epochs: 10
+    verbose: False
+  scheduler_kwargs:
+  #  monitor: &monitor qm9/mae/train
+  #  mode: min
+  #  frequency: 1
+  target_nan_mask: ignore # null: no mask, 0: 0 mask, ignore-flatten, ignore-mean-per-label
+  multitask_handling: flatten # flatten, mean-per-label
+
+# Task-specific
+metrics:
+  pcqm20k_g13: &pcqm_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
+  pcqm20k_n4: *pcqm_metrics
+
+trainer:
+  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/
+    filename: *name
+    #monitor: *monitor
+    #mode: *mode
+    save_top_k: 1
+    every_n_epochs: 100
+  trainer:
+    max_epochs: *max_epochs
+    min_epochs: 1
+    check_val_every_n_epoch: 20
diff --git a/expts/neurips2023_configs/config_small_gated_gcn.yaml b/expts/neurips2023_configs/config_small_gated_gcn.yaml
new file mode 100644
index 000000000..08bbdb4d3
--- /dev/null
+++ b/expts/neurips2023_configs/config_small_gated_gcn.yaml
@@ -0,0 +1,342 @@
+# Testing the gated_gcn model with the PCQMv2 dataset on IPU.
+constants:
+  name: &name neurips2023_small_data_gated_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: 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(5) # IPU would require large batches to be ready for the model.
+    - replicationFactor(8)
+    # - 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"
+  # module_type: "FakeDataModule"  # Option to use generated data
+  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: data/neurips2023/small-dataset/qm9.csv.gz
+        # wget https://storage.googleapis.com/goli-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
+        splits_path: data/neurips2023/small-dataset/qm9_random_splits.pt  # Download with `wget https://storage.googleapis.com/goli-public/datasets/neurips_2023/Small-dataset/qm9_random_splits.pt`
+        seed: *seed
+        task_level: graph
+        label_normalization:
+          method: "normal"
+
+      tox21:
+        df: null
+        df_path: data/neurips2023/small-dataset/Tox21-7k-12-labels.csv.gz
+        # wget https://storage.googleapis.com/goli-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
+        splits_path: data/neurips2023/small-dataset/Tox21_random_splits.pt  # Download with `wget https://storage.googleapis.com/goli-public/datasets/neurips_2023/Small-dataset/Tox21_random_splits.pt`
+        seed: *seed
+        task_level: graph
+
+      zinc:
+        df: null
+        df_path: data/neurips2023/small-dataset/ZINC12k.csv.gz
+        # wget https://storage.googleapis.com/goli-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
+        splits_path: data/neurips2023/small-dataset/ZINC12k_random_splits.pt  # Download with `wget https://storage.googleapis.com/goli-public/datasets/neurips_2023/Small-dataset/ZINC12k_random_splits.pt`
+        seed: *seed
+        task_level: graph
+        label_normalization:
+          method: "normal"
+
+    # Featurization
+    prepare_dict_or_graph: pyg:graph
+    featurization_n_jobs: 30
+    featurization_progress: True
+    featurization_backend: "loky"
+    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
+
+    # cache_data_path: .
+    num_workers: 30 # -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.
+    featurization_backend: "loky"
+
+
+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:   # Set as null to avoid a pre-nn network
+    out_dim: 32
+    hidden_dims: 128
+    depth: 2
+    activation: relu
+    last_activation: none
+    dropout: *dropout
+    normalization: *normalization
+    last_normalization: *normalization
+    residual_type: none
+
+  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
+    out_dim: &gnn_dim 64
+    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:gated-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: 128
+      depth: 2
+      activation: relu
+      last_activation: none
+      dropout: *dropout
+      normalization: *normalization
+      last_normalization: "none"
+      residual_type: none
+    tox21:
+      task_level: graph
+      out_dim: 12
+      hidden_dims: 64
+      depth: 2
+      activation: relu
+      last_activation: sigmoid
+      dropout: *dropout
+      normalization: *normalization
+      last_normalization: "none"
+      residual_type: none
+    zinc:
+      task_level: graph
+      out_dim: 3
+      hidden_dims: 32
+      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_flat"]
+    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 100
+    warmup_epochs: 10
+    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
+  tox21:
+    - name: auroc_flat
+      metric: auroc_ipu
+      multitask_handling: flatten
+      threshold_kwargs: null
+    - name: avpr_flat
+      metric: averageprecision
+      multitask_handling: flatten
+      threshold_kwargs: null
+    - name: f1 > 0.5
+      metric: f1
+      multitask_handling: flatten
+      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: flatten
+      average: micro
+      threshold_kwargs: *threshold_05
+  zinc: *qm9_metrics
+
+trainer:
+  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/
+    filename: *name
+    #monitor: *monitor
+    #mode: *mode
+    save_top_k: 1
+    every_n_epochs: 100
+  trainer:
+    max_epochs: *max_epochs
+    min_epochs: 1
+    check_val_every_n_epoch: 20
diff --git a/expts/neurips2023_configs/config_small_gcn.yaml b/expts/neurips2023_configs/config_small_gcn.yaml
new file mode 100644
index 000000000..3c0a23e27
--- /dev/null
+++ b/expts/neurips2023_configs/config_small_gcn.yaml
@@ -0,0 +1,323 @@
+# 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: 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(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)
+
+datamodule:
+  module_type: "MultitaskFromSmilesDataModule"
+  # module_type: "FakeDataModule"  # Option to use generated data
+  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: data/neurips2023/small-dataset/qm9.csv.gz
+        # wget https://storage.googleapis.com/goli-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
+        splits_path: data/neurips2023/small-dataset/qm9_random_splits.pt  # Download with `wget https://storage.googleapis.com/goli-public/datasets/neurips_2023/Small-dataset/qm9_random_splits.pt`
+        seed: *seed
+        task_level: graph
+        label_normalization:
+          method: "normal"
+
+      tox21:
+        df: null
+        df_path: data/neurips2023/small-dataset/Tox21-7k-12-labels.csv.gz
+        # wget https://storage.googleapis.com/goli-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
+        splits_path: data/neurips2023/small-dataset/Tox21_random_splits.pt  # Download with `wget https://storage.googleapis.com/goli-public/datasets/neurips_2023/Small-dataset/Tox21_random_splits.pt`
+        seed: *seed
+        task_level: graph
+
+      zinc:
+        df: null
+        df_path: data/neurips2023/small-dataset/ZINC12k.csv.gz
+        # wget https://storage.googleapis.com/goli-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
+        splits_path: data/neurips2023/small-dataset/ZINC12k_random_splits.pt  # Download with `wget https://storage.googleapis.com/goli-public/datasets/neurips_2023/Small-dataset/ZINC12k_random_splits.pt`
+        seed: *seed
+        task_level: graph
+        label_normalization:
+          method: "normal"
+
+    # Featurization
+    prepare_dict_or_graph: pyg:graph
+    featurization_n_jobs: 30
+    featurization_progress: True
+    featurization_backend: "loky"
+    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
+
+    # cache_data_path: .
+    num_workers: 30 # -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.
+    featurization_backend: "loky"
+
+
+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 128
+    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:
+    qm9:
+      task_level: graph
+      out_dim: 19
+      hidden_dims: 128
+      depth: 2
+      activation: relu
+      last_activation: none
+      dropout: *dropout
+      normalization: *normalization
+      last_normalization: "none"
+      residual_type: none
+    tox21:
+      task_level: graph
+      out_dim: 12
+      hidden_dims: 64
+      depth: 2
+      activation: relu
+      last_activation: sigmoid
+      dropout: *dropout
+      normalization: *normalization
+      last_normalization: "none"
+      residual_type: none
+    zinc:
+      task_level: graph
+      out_dim: 3
+      hidden_dims: 32
+      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_flat"]
+    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 100
+    warmup_epochs: 10
+    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
+  tox21:
+    - name: auroc_flat
+      metric: auroc_ipu
+      multitask_handling: flatten
+      threshold_kwargs: null
+    - name: avpr_flat
+      metric: averageprecision
+      multitask_handling: flatten
+      threshold_kwargs: null
+    - name: f1 > 0.5
+      metric: f1
+      multitask_handling: flatten
+      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: flatten
+      average: micro
+      threshold_kwargs: *threshold_05
+  zinc: *qm9_metrics
+
+trainer:
+  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/
+    filename: *name
+    #monitor: *monitor
+    #mode: *mode
+    save_top_k: 1
+    every_n_epochs: 100
+  trainer:
+    max_epochs: *max_epochs
+    min_epochs: 1
+    check_val_every_n_epoch: 20
diff --git a/expts/neurips2023_configs/config_small_gin.yaml b/expts/neurips2023_configs/config_small_gin.yaml
new file mode 100644
index 000000000..8cedb9b8e
--- /dev/null
+++ b/expts/neurips2023_configs/config_small_gin.yaml
@@ -0,0 +1,334 @@
+# Testing the gin model with the PCQMv2 dataset on IPU.
+constants:
+  name: &name neurips2023_small_data_gin
+  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(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"
+  # module_type: "FakeDataModule"  # Option to use generated data
+  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: data/neurips2023/small-dataset/qm9.csv.gz
+        # wget https://storage.googleapis.com/goli-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
+        splits_path: data/neurips2023/small-dataset/qm9_random_splits.pt  # Download with `wget https://storage.googleapis.com/goli-public/datasets/neurips_2023/Small-dataset/qm9_random_splits.pt`
+        seed: *seed
+        task_level: graph
+        label_normalization:
+          method: "normal"
+
+      tox21:
+        df: null
+        df_path: data/neurips2023/small-dataset/Tox21-7k-12-labels.csv.gz
+        # wget https://storage.googleapis.com/goli-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
+        splits_path: data/neurips2023/small-dataset/Tox21_random_splits.pt  # Download with `wget https://storage.googleapis.com/goli-public/datasets/neurips_2023/Small-dataset/Tox21_random_splits.pt`
+        seed: *seed
+        task_level: graph
+
+      zinc:
+        df: null
+        df_path: data/neurips2023/small-dataset/ZINC12k.csv.gz
+        # wget https://storage.googleapis.com/goli-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
+        splits_path: data/neurips2023/small-dataset/ZINC12k_random_splits.pt  # Download with `wget https://storage.googleapis.com/goli-public/datasets/neurips_2023/Small-dataset/ZINC12k_random_splits.pt`
+        seed: *seed
+        task_level: graph
+        label_normalization:
+          method: "normal"
+
+    # Featurization
+    prepare_dict_or_graph: pyg:graph
+    featurization_n_jobs: 30
+    featurization_progress: True
+    featurization_backend: "loky"
+    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
+
+    # cache_data_path: .
+    num_workers: 30 # -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.
+    featurization_backend: "loky"
+
+
+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 64
+    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:gin' #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: 128
+      depth: 2
+      activation: relu
+      last_activation: none
+      dropout: *dropout
+      normalization: *normalization
+      last_normalization: "none"
+      residual_type: none
+    tox21:
+      task_level: graph
+      out_dim: 12
+      hidden_dims: 64
+      depth: 2
+      activation: relu
+      last_activation: sigmoid
+      dropout: *dropout
+      normalization: *normalization
+      last_normalization: "none"
+      residual_type: none
+    zinc:
+      task_level: graph
+      out_dim: 3
+      hidden_dims: 32
+      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_flat"]
+    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 100
+    warmup_epochs: 10
+    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
+  tox21:
+    - name: auroc_flat
+      metric: auroc_ipu
+      multitask_handling: flatten
+      threshold_kwargs: null
+    - name: avpr_flat
+      metric: averageprecision
+      multitask_handling: flatten
+      threshold_kwargs: null
+    - name: f1 > 0.5
+      metric: f1
+      multitask_handling: flatten
+      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: flatten
+      average: micro
+      threshold_kwargs: *threshold_05
+  zinc: *qm9_metrics
+
+trainer:
+  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/
+    filename: *name
+    #monitor: *monitor
+    #mode: *mode
+    save_top_k: 1
+    every_n_epochs: 100
+  trainer:
+    max_epochs: *max_epochs
+    min_epochs: 1
+    check_val_every_n_epoch: 20
diff --git a/expts/neurips2023_configs/config_small_gine.yaml b/expts/neurips2023_configs/config_small_gine.yaml
new file mode 100644
index 000000000..78be87b34
--- /dev/null
+++ b/expts/neurips2023_configs/config_small_gine.yaml
@@ -0,0 +1,342 @@
+# Testing the gine model with the PCQMv2 dataset on IPU.
+constants:
+  name: &name neurips2023_small_data_gine
+  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(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"
+  # module_type: "FakeDataModule"  # Option to use generated data
+  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: data/neurips2023/small-dataset/qm9.csv.gz
+        # wget https://storage.googleapis.com/goli-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
+        splits_path: data/neurips2023/small-dataset/qm9_random_splits.pt  # Download with `wget https://storage.googleapis.com/goli-public/datasets/neurips_2023/Small-dataset/qm9_random_splits.pt`
+        seed: *seed
+        task_level: graph
+        label_normalization:
+          method: "normal"
+
+      tox21:
+        df: null
+        df_path: data/neurips2023/small-dataset/Tox21-7k-12-labels.csv.gz
+        # wget https://storage.googleapis.com/goli-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
+        splits_path: data/neurips2023/small-dataset/Tox21_random_splits.pt  # Download with `wget https://storage.googleapis.com/goli-public/datasets/neurips_2023/Small-dataset/Tox21_random_splits.pt`
+        seed: *seed
+        task_level: graph
+
+      zinc:
+        df: null
+        df_path: data/neurips2023/small-dataset/ZINC12k.csv.gz
+        # wget https://storage.googleapis.com/goli-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
+        splits_path: data/neurips2023/small-dataset/ZINC12k_random_splits.pt  # Download with `wget https://storage.googleapis.com/goli-public/datasets/neurips_2023/Small-dataset/ZINC12k_random_splits.pt`
+        seed: *seed
+        task_level: graph
+        label_normalization:
+          method: "normal"
+
+    # Featurization
+    prepare_dict_or_graph: pyg:graph
+    featurization_n_jobs: 30
+    featurization_progress: True
+    featurization_backend: "loky"
+    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
+
+    # cache_data_path: .
+    num_workers: 30 # -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.
+    featurization_backend: "loky"
+
+
+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:   # Set as null to avoid a pre-nn network
+    out_dim: 32
+    hidden_dims: 128
+    depth: 2
+    activation: relu
+    last_activation: none
+    dropout: *dropout
+    normalization: *normalization
+    last_normalization: *normalization
+    residual_type: none
+
+  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
+    out_dim: &gnn_dim 64
+    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:gine' #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: 128
+      depth: 2
+      activation: relu
+      last_activation: none
+      dropout: *dropout
+      normalization: *normalization
+      last_normalization: "none"
+      residual_type: none
+    tox21:
+      task_level: graph
+      out_dim: 12
+      hidden_dims: 64
+      depth: 2
+      activation: relu
+      last_activation: sigmoid
+      dropout: *dropout
+      normalization: *normalization
+      last_normalization: "none"
+      residual_type: none
+    zinc:
+      task_level: graph
+      out_dim: 3
+      hidden_dims: 32
+      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_flat"]
+    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 100
+    warmup_epochs: 10
+    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
+  tox21:
+    - name: auroc_flat
+      metric: auroc_ipu
+      multitask_handling: flatten
+      threshold_kwargs: null
+    - name: avpr_flat
+      metric: averageprecision
+      multitask_handling: flatten
+      threshold_kwargs: null
+    - name: f1 > 0.5
+      metric: f1
+      multitask_handling: flatten
+      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: flatten
+      average: micro
+      threshold_kwargs: *threshold_05
+  zinc: *qm9_metrics
+
+trainer:
+  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/
+    filename: *name
+    #monitor: *monitor
+    #mode: *mode
+    save_top_k: 1
+    every_n_epochs: 100
+  trainer:
+    max_epochs: *max_epochs
+    min_epochs: 1
+    check_val_every_n_epoch: 20
diff --git a/expts/neurips2023_configs/config_small_mpnn.yaml b/expts/neurips2023_configs/config_small_mpnn.yaml
new file mode 100644
index 000000000..a94200f51
--- /dev/null
+++ b/expts/neurips2023_configs/config_small_mpnn.yaml
@@ -0,0 +1,351 @@
+# Testing the mpnn only model with the PCQMv2 dataset on IPU.
+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(8)
+#     # - 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"
+  # module_type: "FakeDataModule"  # Option to use generated data
+  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: data/neurips2023/small-dataset/qm9.csv.gz
+        # wget https://storage.googleapis.com/goli-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
+        splits_path: data/neurips2023/small-dataset/qm9_random_splits.pt  # Download with `wget https://storage.googleapis.com/goli-public/datasets/neurips_2023/Small-dataset/qm9_random_splits.pt`
+        seed: *seed
+        task_level: graph
+        label_normalization:
+          method: "normal"
+
+      tox21:
+        df: null
+        df_path: data/neurips2023/small-dataset/Tox21-7k-12-labels.csv.gz
+        # wget https://storage.googleapis.com/goli-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
+        splits_path: data/neurips2023/small-dataset/Tox21_random_splits.pt  # Download with `wget https://storage.googleapis.com/goli-public/datasets/neurips_2023/Small-dataset/Tox21_random_splits.pt`
+        seed: *seed
+        task_level: graph
+
+      zinc:
+        df: null
+        df_path: data/neurips2023/small-dataset/ZINC12k.csv.gz
+        # wget https://storage.googleapis.com/goli-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
+        splits_path: data/neurips2023/small-dataset/ZINC12k_random_splits.pt  # Download with `wget https://storage.googleapis.com/goli-public/datasets/neurips_2023/Small-dataset/ZINC12k_random_splits.pt`
+        seed: *seed
+        task_level: graph
+        label_normalization:
+          method: "normal"
+
+    # Featurization
+    prepare_dict_or_graph: pyg:graph
+    featurization_n_jobs: 30
+    featurization_progress: True
+    featurization_backend: "loky"
+    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
+
+    # cache_data_path: .
+    num_workers: 30 # -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.
+    featurization_backend: "loky"
+
+
+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:   # Set as null to avoid a pre-nn network
+    out_dim: 32
+    hidden_dims: 128
+    depth: 2
+    activation: relu
+    last_activation: none
+    dropout: *dropout
+    normalization: *normalization
+    last_normalization: *normalization
+    residual_type: none
+
+  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
+    out_dim: &gnn_dim 64
+    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:gps' #pyg:gine #'pyg:gps' # pyg:gated-gcn, pyg:gine,pyg:gps
+    layer_kwargs:  # Parameters for the model itself. You could define dropout_attn: 0.1
+      node_residual: false
+      mpnn_type: 'pyg:mpnnplus'
+      mpnn_kwargs:
+        in_dim: *gnn_dim
+        out_dim: *gnn_dim
+        in_dim_edges: 32
+        out_dim_edges: 32
+      attn_type: "none" # "full-attention", "none"
+      attn_kwargs: null
+
+
+  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: 128
+      depth: 2
+      activation: relu
+      last_activation: none
+      dropout: *dropout
+      normalization: *normalization
+      last_normalization: "none"
+      residual_type: none
+    tox21:
+      task_level: graph
+      out_dim: 12
+      hidden_dims: 64
+      depth: 2
+      activation: relu
+      last_activation: sigmoid
+      dropout: *dropout
+      normalization: *normalization
+      last_normalization: "none"
+      residual_type: none
+    zinc:
+      task_level: graph
+      out_dim: 3
+      hidden_dims: 32
+      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_flat"]
+    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 100
+    warmup_epochs: 10
+    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
+  tox21:
+    - name: auroc_flat
+      metric: auroc_ipu
+      multitask_handling: flatten
+      threshold_kwargs: null
+    - name: avpr_flat
+      metric: averageprecision
+      multitask_handling: flatten
+      threshold_kwargs: null
+    - name: f1 > 0.5
+      metric: f1
+      multitask_handling: flatten
+      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: flatten
+      average: micro
+      threshold_kwargs: *threshold_05
+  zinc: *qm9_metrics
+
+trainer:
+  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/
+    filename: *name
+    #monitor: *monitor
+    #mode: *mode
+    save_top_k: 1
+    every_n_epochs: 100
+  trainer:
+    max_epochs: *max_epochs
+    min_epochs: 1
+    check_val_every_n_epoch: 20
diff --git a/goli/config/_loader.py b/goli/config/_loader.py
index fb753880a..7ddc742f8 100644
--- a/goli/config/_loader.py
+++ b/goli/config/_loader.py
@@ -132,7 +132,6 @@ def load_datamodule(
             gradient_accumulation=config["trainer"]["trainer"].get("accumulate_grad_batches", None),
             ipu_inference_opts=ipu_inference_opts,
         )
-
         # Define the Dataloader options for the IPU on the training sets
         bz_train = cfg_data["batch_size_training"]
         ipu_dataloader_training_opts = IPUDataloaderOptions(
@@ -180,7 +179,6 @@ def load_metrics(config: Union[omegaconf.DictConfig, Dict[str, Any]]) -> Dict[st
         for this_metric in cfg_metrics[task]:
             name = this_metric.pop("name")
             task_metrics[task][name] = MetricWrapper(**this_metric)
-
     return task_metrics
 
 
diff --git a/goli/config/fake_and_missing_multilevel_multitask_pyg.yaml b/goli/config/fake_and_missing_multilevel_multitask_pyg.yaml
new file mode 100644
index 000000000..71d2ef1e2
--- /dev/null
+++ b/goli/config/fake_and_missing_multilevel_multitask_pyg.yaml
@@ -0,0 +1,184 @@
+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"
+      score:
+        df: null
+        task_level: "edge"
+        df_path: "./tests/fake_and_missing_multilevel_data.parquet"
+        smiles_col: "ordered_smiles"
+        label_cols: ["edge_label_list", "edge_label_np"]
+        split_val: 0.2
+        split_test: 0.2
+        seed: 19
+        splits_path: null
+        sample_size: null
+        idx_col: null
+        weights_col: null
+        weights_type: null
+      logp:
+        df: null
+        task_level: "node"
+        df_path: "./tests/fake_and_missing_multilevel_data.parquet"
+        smiles_col: "ordered_smiles"
+        label_cols: ["node_label_list", "node_label_np"]
+        split_val: 0.2
+        split_test: 0.2
+        seed: 19
+        splits_path: null
+        sample_size: null
+        idx_col: null
+        weights_col: null
+        weights_type: null
+      SA:
+        df: null
+        task_level: "graph"
+        df_path: "./tests/fake_and_missing_multilevel_data.parquet"
+        smiles_col: "ordered_smiles"
+        label_cols: ["graph_label"]
+        split_val: 0.2
+        split_test: 0.2
+        seed: 19
+        splits_path: null                 # This may not always be provided
+        sample_size: null                 # This may not always be provided
+        idx_col: null                     # This may not always be provided
+        weights_col: null                 # This may not always be provided
+
+    # Featurization
+    featurization_n_jobs: 16
+    featurization_progress: True
+    featurization:
+      atom_property_list_onehot: ["atomic-number", "degree"]
+      atom_property_list_float: []
+      edge_property_list: ["in-ring", "bond-type-onehot"]
+      add_self_loop: False
+      explicit_H: False
+      use_bonds_weights: False
+
+    # Data handling-related
+    batch_size_training: 16
+    batch_size_inference: 16
+    # cache_data_path: null
+
+architecture:     # The parameters for the full graph network are taken from `config_micro_ZINC.yaml`
+  model_type: FullGraphMultiTaskNetwork
+  pre_nn:         # Set as null to avoid a pre-nn network
+    out_dim: 32
+    hidden_dims: 32
+    depth: 1
+    activation: relu
+    last_activation: none
+    dropout: &dropout 0.1
+    normalization: &normalization "batch_norm"
+    last_normalization: *normalization
+    residual_type: none
+
+  pre_nn_edges:   # Set as null to avoid a pre-nn network
+    out_dim: 16
+    hidden_dims: 16
+    depth: 2
+    activation: relu
+    last_activation: none
+    dropout: *dropout
+    normalization: *normalization
+    last_normalization: *normalization
+    residual_type: none
+
+  gnn:            # Set as null to avoid a post-nn network
+    out_dim: 32
+    hidden_dims: 32
+    depth: 4
+    activation: relu
+    last_activation: none
+    dropout: *dropout
+    normalization: *normalization
+    last_normalization: *normalization
+    residual_type: random
+    virtual_node: 'sum'
+    layer_type: 'pyg:pna-msgpass'
+    layer_kwargs:
+      # num_heads: 3
+      aggregators: [mean, max]
+      scalers: [identity, amplification, attenuation]
+
+  graph_output_nn:
+    node:
+      out_dim: 16
+      hidden_dims: 32
+      depth: 2
+      activation: relu
+      last_activation: none
+      dropout: *dropout
+      normalization: *normalization
+      last_normalization: "none"
+      residual_type: none 
+    graph:
+      pooling: [sum, max]
+      out_dim: 1
+      hidden_dims: 32
+      depth: 2
+      activation: relu
+      last_activation: none
+      dropout: *dropout
+      normalization: *normalization
+      last_normalization: "none"
+      residual_type: none 
+    edge:
+      out_dim: 16
+      hidden_dims: 32
+      depth: 2
+      activation: relu
+      last_activation: none
+      dropout: *dropout
+      normalization: *normalization
+      last_normalization: "none"
+      residual_type: none 
+    nodepair:
+      out_dim: 16
+      hidden_dims: 32
+      depth: 2
+      activation: relu
+      last_activation: none
+      dropout: *dropout
+      normalization: *normalization
+      last_normalization: "none"
+      residual_type: none     
+
+  task_heads:     # Set as null to avoid task heads. Recall that the arguments for the TaskHeads is a List of TaskHeadParams
+    task_1:
+      task_level: "node"
+      out_dim: 5
+      hidden_dims: [5, 6, 7]
+      #depth: none                          # Not needed if we have hidden_dims
+      activation: relu
+      last_activation: none
+      dropout: 0.2
+      normalization: none
+      residual_type: none
+    task_2:
+      task_level: "edge"
+      out_dim: 3
+      hidden_dims: [8, 9, 10]
+      activation: relu
+      last_activation: none
+      dropout: 0.2
+      normalization: none
+      residual_type: none
+    task_3:
+      task_level: "graph"
+      out_dim: 4
+      hidden_dims: [2, 2, 2]
+      activation: relu
+      last_activation: none
+      dropout: 0.2
+      normalization: none
+      residual_type: none
+    task_4:
+      task_level: "nodepair"
+      out_dim: 4
+      hidden_dims: [2, 2, 2]
+      activation: relu
+      last_activation: none
+      dropout: 0.2
+      normalization: none
+      residual_type: none
\ No newline at end of file
diff --git a/goli/config/fake_multilevel_multitask_pyg.yaml b/goli/config/fake_multilevel_multitask_pyg.yaml
new file mode 100644
index 000000000..9fbe7f5ac
--- /dev/null
+++ b/goli/config/fake_multilevel_multitask_pyg.yaml
@@ -0,0 +1,184 @@
+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"
+      score:
+        df: null
+        task_level: "edge"
+        df_path: "./tests/converted_fake_multilevel_data.parquet"
+        smiles_col: "ordered_smiles"
+        label_cols: ["edge_label_list", "edge_label_np"]
+        split_val: 0.2
+        split_test: 0.2
+        seed: 19
+        splits_path: null
+        sample_size: null
+        idx_col: null
+        weights_col: null
+        weights_type: null
+      logp:
+        df: null
+        task_level: "node"
+        df_path: "./tests/converted_fake_multilevel_data.parquet"
+        smiles_col: "ordered_smiles"
+        label_cols: ["node_label_list", "node_label_np"]
+        split_val: 0.2
+        split_test: 0.2
+        seed: 19
+        splits_path: null
+        sample_size: null
+        idx_col: null
+        weights_col: null
+        weights_type: null
+      SA:
+        df: null
+        task_level: "graph"
+        df_path: "./tests/converted_fake_multilevel_data.parquet"
+        smiles_col: "ordered_smiles"
+        label_cols: ["graph_label"]
+        split_val: 0.2
+        split_test: 0.2
+        seed: 19
+        splits_path: null                 # This may not always be provided
+        sample_size: null                 # This may not always be provided
+        idx_col: null                     # This may not always be provided
+        weights_col: null                 # This may not always be provided
+
+    # Featurization
+    featurization_n_jobs: 16
+    featurization_progress: True
+    featurization:
+      atom_property_list_onehot: ["atomic-number", "degree"]
+      atom_property_list_float: []
+      edge_property_list: ["in-ring", "bond-type-onehot"]
+      add_self_loop: False
+      explicit_H: False
+      use_bonds_weights: False
+
+    # Data handling-related
+    batch_size_training: 16
+    batch_size_inference: 16
+    # cache_data_path: null
+
+architecture:     # The parameters for the full graph network are taken from `config_micro_ZINC.yaml`
+  model_type: FullGraphMultiTaskNetwork
+  pre_nn:         # Set as null to avoid a pre-nn network
+    out_dim: 32
+    hidden_dims: 32
+    depth: 1
+    activation: relu
+    last_activation: none
+    dropout: &dropout 0.1
+    normalization: &normalization "batch_norm"
+    last_normalization: *normalization
+    residual_type: none
+
+  pre_nn_edges:   # Set as null to avoid a pre-nn network
+    out_dim: 16
+    hidden_dims: 16
+    depth: 2
+    activation: relu
+    last_activation: none
+    dropout: *dropout
+    normalization: *normalization
+    last_normalization: *normalization
+    residual_type: none
+
+  gnn:            # Set as null to avoid a post-nn network
+    out_dim: 32
+    hidden_dims: 32
+    depth: 4
+    activation: relu
+    last_activation: none
+    dropout: *dropout
+    normalization: *normalization
+    last_normalization: *normalization
+    residual_type: random
+    virtual_node: 'sum'
+    layer_type: 'pyg:pna-msgpass'
+    layer_kwargs:
+      # num_heads: 3
+      aggregators: [mean, max]
+      scalers: [identity, amplification, attenuation]
+
+  graph_output_nn:
+    node:
+      out_dim: 16
+      hidden_dims: 32
+      depth: 2
+      activation: relu
+      last_activation: none
+      dropout: *dropout
+      normalization: *normalization
+      last_normalization: "none"
+      residual_type: none 
+    graph:
+      pooling: [sum, max]
+      out_dim: 1
+      hidden_dims: 32
+      depth: 2
+      activation: relu
+      last_activation: none
+      dropout: *dropout
+      normalization: *normalization
+      last_normalization: "none"
+      residual_type: none 
+    edge:
+      out_dim: 16
+      hidden_dims: 32
+      depth: 2
+      activation: relu
+      last_activation: none
+      dropout: *dropout
+      normalization: *normalization
+      last_normalization: "none"
+      residual_type: none 
+    nodepair:
+      out_dim: 16
+      hidden_dims: 32
+      depth: 2
+      activation: relu
+      last_activation: none
+      dropout: *dropout
+      normalization: *normalization
+      last_normalization: "none"
+      residual_type: none     
+
+  task_heads:     # Set as null to avoid task heads. Recall that the arguments for the TaskHeads is a List of TaskHeadParams
+    task_1:
+      task_level: "node"
+      out_dim: 5
+      hidden_dims: [5, 6, 7]
+      #depth: none                          # Not needed if we have hidden_dims
+      activation: relu
+      last_activation: none
+      dropout: 0.2
+      normalization: none
+      residual_type: none
+    task_2:
+      task_level: "edge"
+      out_dim: 3
+      hidden_dims: [8, 9, 10]
+      activation: relu
+      last_activation: none
+      dropout: 0.2
+      normalization: none
+      residual_type: none
+    task_3:
+      task_level: "graph"
+      out_dim: 4
+      hidden_dims: [2, 2, 2]
+      activation: relu
+      last_activation: none
+      dropout: 0.2
+      normalization: none
+      residual_type: none
+    task_4:
+      task_level: "nodepair"
+      out_dim: 4
+      hidden_dims: [2, 2, 2]
+      activation: relu
+      last_activation: none
+      dropout: 0.2
+      normalization: none
+      residual_type: none
\ No newline at end of file
diff --git a/goli/config/zinc_default_multitask_pyg.yaml b/goli/config/zinc_default_multitask_pyg.yaml
index 63459abbe..2d3d86775 100644
--- a/goli/config/zinc_default_multitask_pyg.yaml
+++ b/goli/config/zinc_default_multitask_pyg.yaml
@@ -4,6 +4,7 @@ datamodule:
     task_specific_args:   # To be replaced by a new class "DatasetParams"
       SA:
         df: null
+        task_level: "graph"
         df_path: "goli/data/multitask/tiny_ZINC_SA.csv"
         smiles_col: "SMILES"
         label_cols: ["SA"]
@@ -16,6 +17,7 @@ datamodule:
         weights_col: null                 # This may not always be provided
       logp:
         df: null
+        task_level: "graph"
         df_path: "goli/data/multitask/tiny_ZINC_logp.csv"
         smiles_col: "SMILES"
         label_cols: ["logp"]
@@ -29,6 +31,7 @@ datamodule:
         weights_type: null
       score:
         df: null
+        task_level: "graph"
         df_path: "goli/data/multitask/tiny_ZINC_score.csv"
         smiles_col: "SMILES"
         label_cols: ["score"]
diff --git a/goli/data/collate.py b/goli/data/collate.py
index bb025d240..6514eb819 100644
--- a/goli/data/collate.py
+++ b/goli/data/collate.py
@@ -10,6 +10,7 @@
 
 from goli.features import GraphDict, to_dense_array
 from goli.utils.packing import fast_packing, get_pack_sizes, node_to_pack_indices_mask
+from loguru import logger
 
 
 def goli_collate_fn(
@@ -65,13 +66,17 @@ def goli_collate_fn(
     """
 
     elem = elements[0]
-
     if isinstance(elem, Mapping):
         batch = {}
         for key in elem:
+            # Multitask setting: We have to pad the missing labels
+            if key == "labels":
+                labels = [d[key] for d in elements]
+                batch[key] = collate_labels(labels, labels_size_dict)
+
             # If the features are a dictionary containing GraphDict elements,
             # Convert to pyg graphs and use the pyg batching.
-            if isinstance(elem[key], GraphDict):
+            elif isinstance(elem[key], GraphDict):
                 pyg_graphs = [d[key].make_pyg_graph(mask_nan=mask_nan) for d in elements]
                 batch[key] = collage_pyg_graph(pyg_graphs)
 
@@ -83,12 +88,6 @@ def goli_collate_fn(
             # Ignore the collate for specific keys
             elif key in do_not_collate_keys:
                 batch[key] = [d[key] for d in elements]
-
-            # Multitask setting: We have to pad the missing labels
-            elif key == "labels":
-                labels = [d[key] for d in elements]
-                batch[key] = collate_labels(labels, labels_size_dict)
-
             # Otherwise, use the default torch batching
             else:
                 batch[key] = default_collate([d[key] for d in elements])
@@ -162,8 +161,67 @@ def collage_pyg_graph(pyg_graphs: Iterable[Union[Data, Dict]], batch_size_per_pa
     return Batch.from_data_list(pyg_batch)
 
 
+def pad_to_expected_label_size(labels: torch.Tensor, label_size: List[int]):
+    """Determine difference of ``labels`` shape to expected shape `label_size` and pad
+    with ``torch.nan`` accordingly.
+    """
+    if label_size == list(labels.shape):
+        return labels
+
+    missing_dims = len(label_size) - len(labels.shape)
+    for _ in range(missing_dims):
+        labels.unsqueeze(-1)
+
+    pad_sizes = [(0, expected - actual) for expected, actual in zip(label_size, labels.shape)]
+    pad_sizes = [item for before_after in pad_sizes for item in before_after]
+    pad_sizes.reverse()
+
+    if any([s < 0 for s in pad_sizes]):
+        logger.warning(f"More labels available than expected. Will remove data to fit expected size.")
+
+    return torch.nn.functional.pad(labels, pad_sizes, value=torch.nan)
+
+
+def collate_pyg_graph_labels(pyg_labels: List[Data]):
+    """
+    Function to collate pytorch geometric labels.
+    Convert all numpy types to torch
+
+    Parameters:
+        pyg_labels: Iterable of PyG label Data objects
+    """
+    pyg_batch = []
+    for pyg_label in pyg_labels:
+        for pyg_key in set(pyg_label.keys) - set(["x", "edge_index"]):
+            tensor = pyg_label[pyg_key]
+            # Convert numpy/scipy to Pytorch
+            if isinstance(tensor, (ndarray, spmatrix)):
+                tensor = torch.as_tensor(to_dense_array(tensor, tensor.dtype))
+
+            pyg_label[pyg_key] = tensor
+
+        pyg_batch.append(pyg_label)
+
+    return Batch.from_data_list(pyg_batch)
+
+
+def get_expected_label_size(label_data: Data, task: str, label_size: List[int]):
+    """Determines expected label size based on the specfic graph properties
+    and the number of targets in the task-dataset.
+    """
+    if task.startswith("graph_"):
+        num_labels = 1
+    elif task.startswith("node_"):
+        num_labels = label_data.x.size(0)
+    elif task.startswith("edge_"):
+        num_labels = label_data.edge_index.size(1)
+    elif task.startswith("nodepair_"):
+        raise NotImplementedError()
+    return [num_labels] + label_size
+
+
 def collate_labels(
-    labels: List[Dict[str, torch.Tensor]],
+    labels: List[Data],
     labels_size_dict: Optional[Dict[str, Any]] = None,
 ):
     """Collate labels for multitask learning.
@@ -178,19 +236,44 @@ def collate_labels(
         A dictionary of the form Dict[tasks, labels] where tasks is the name of the task and labels
         is a tensor of shape (batch_size, *labels_size_dict[task]).
     """
-    labels_dict = {}
-
     if labels_size_dict is not None:
         for this_label in labels:
-            empty_task_labels = set(labels_size_dict.keys()) - set(this_label.keys())
+            for task in labels_size_dict.keys():
+                labels_size_dict[task] = list(labels_size_dict[task])
+                if len(labels_size_dict[task]) >= 2:
+                    labels_size_dict[task] = labels_size_dict[task][1:]
+                elif not task.startswith("graph_"):
+                    labels_size_dict[task] = [1]
+
+            empty_task_labels = set(labels_size_dict.keys()) - set(this_label.keys)
             for task in empty_task_labels:
+                labels_size_dict[task] = get_expected_label_size(this_label, task, labels_size_dict[task])
                 this_label[task] = torch.full([*labels_size_dict[task]], torch.nan)
-            for task in this_label.keys():
-                if not isinstance(task, torch.Tensor):
+
+            for task in set(this_label.keys) - set(["x", "edge_index"]) - empty_task_labels:
+                labels_size_dict[task] = get_expected_label_size(this_label, task, labels_size_dict[task])
+
+                if not isinstance(this_label[task], (torch.Tensor)):
                     this_label[task] = torch.as_tensor(this_label[task])
-    labels_dict = default_collate(labels)
 
-    return labels_dict
+                # Ensure explicit task dimension also for single task labels
+                if len(this_label[task].shape) == 1:
+                    # Distinguish whether target dim or entity dim is missing
+                    if labels_size_dict[task][0] == this_label[task].shape[0]:
+                        # num graphs/nodes/edges/nodepairs already matching
+                        this_label[task] = this_label[task].unsqueeze(1)
+                    else:
+                        # data lost unless entity dim is supposed to be 1
+                        if labels_size_dict[task][0] == 1:
+                            this_label[task] = this_label[task].unsqueeze(0)
+                        else:
+                            raise ValueError(
+                                f"Labels for {labels_size_dict[task][0]} nodes/edges/nodepairs expected, got 1."
+                            )
+
+                this_label[task] = pad_to_expected_label_size(this_label[task], labels_size_dict[task])
+
+    return collate_pyg_graph_labels(labels)
 
 
 def pad_nodepairs(pe: torch.Tensor, num_nodes: int, max_num_nodes_per_graph: int):
diff --git a/goli/data/datamodule.py b/goli/data/datamodule.py
index 3b6358877..c227f9f2a 100644
--- a/goli/data/datamodule.py
+++ b/goli/data/datamodule.py
@@ -8,6 +8,7 @@
 from copy import deepcopy
 import time
 import gc
+from rdkit import Chem
 
 from loguru import logger
 import fsspec
@@ -26,6 +27,7 @@
 from pytorch_lightning.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
 
@@ -46,9 +48,11 @@
 from goli.data.collate import goli_collate_fn
 import goli.data.dataset as Datasets
 from goli.data.normalization import LabelNormalization
+from goli.data.multilevel_utils import extract_labels
 
 torch.multiprocessing.set_sharing_strategy("file_system")
 
+
 PCQM4M_meta = {
     "num tasks": 1,
     "eval metric": "mae",
@@ -342,9 +346,9 @@ def _read_parquet(path, **kwargs):
             first_elem = this_series.values[0]
             is_float = False
             if isinstance(first_elem, (list, tuple)):
-                is_float = isinstance(first_elem[0], np.floating)
+                is_float = first_elem[0].dtype.kind == 'f'
             elif isinstance(first_elem, np.ndarray):
-                is_float = isinstance(first_elem, np.floating)
+                is_float = first_elem.dtype.kind == 'f'
 
             # Convert floats to float16
             if is_float:
@@ -631,6 +635,7 @@ 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,
         df_path: Optional[Union[str, os.PathLike]] = None,
         smiles_col: str = None,
@@ -649,6 +654,7 @@ def __init__(
         """
         object to store the parameters for the dataset processing
         Parameters:
+            task_level: The task level, wether it is graph, node, edge or nodepair
             df: The dataframe containing the data
             df_path: The path to the dataframe containing the data
             smiles_col: The column name of the smiles
@@ -663,6 +669,7 @@ def __init__(
             splits_path: The path to the splits
         """
         self.df = df
+        self.task_level = task_level
         self.df_path = df_path
         self.smiles_col = smiles_col
         self.label_cols = label_cols
@@ -847,7 +854,8 @@ def __init__(
 
         # TODO: Have the input argument to the Data Module be of type DatasetParams
         self.task_dataset_processing_params = {
-            task: DatasetProcessingParams(**ds_args) for task, ds_args in task_specific_args.items()
+            self._get_task_key(ds_args["task_level"], task): DatasetProcessingParams(**ds_args)
+            for task, ds_args in task_specific_args.items()
         }
         self.featurization_n_jobs = featurization_n_jobs
         self.featurization_progress = featurization_progress
@@ -888,6 +896,12 @@ def __init__(
             )
         self.data_hash = self.get_data_hash()
 
+    def _get_task_key(self, task_level: str, task: str):
+        task_prefix = f"{task_level}_"
+        if not task.startswith(task_prefix):
+            task = task_prefix + task
+        return task
+
     def prepare_data(self):
         """Called only from a single process in distributed settings. Steps:
 
@@ -970,6 +984,7 @@ def prepare_data(self):
             args = self.task_dataset_processing_params[task]
             smiles, labels, sample_idx, extras = self._extract_smiles_labels(
                 df,
+                task_level=args.task_level,
                 smiles_col=args.smiles_col,
                 label_cols=args.label_cols,
                 idx_col=args.idx_col,
@@ -1168,12 +1183,14 @@ def _make_multitask_dataset(
         if stage == "train":
             singletask_datasets = self.train_singletask_datasets
             about = "training set"
-        if stage == "val":
+        elif stage == "val":
             singletask_datasets = self.val_singletask_datasets
             about = "validation set"
-        if stage == "test":
+        elif stage == "test":
             singletask_datasets = self.test_singletask_datasets
             about = "test set"
+        else:
+            raise ValueError(f"Unknown stage {stage}")
 
         if load_from_file is None:
             load_from_file = self.load_from_file
@@ -1188,7 +1205,7 @@ def _make_multitask_dataset(
         else:
             files_ready = False
 
-        return Datasets.MultitaskDataset(
+        multitask_dataset = Datasets.MultitaskDataset(
             singletask_datasets,
             n_jobs=self.featurization_n_jobs,
             backend=self.featurization_backend,
@@ -1201,6 +1218,14 @@ def _make_multitask_dataset(
             files_ready=files_ready,
         )  # type: ignore
 
+        if stage == "train":
+            self.get_label_statistics(
+                self.processed_graph_data_path, self.data_hash, multitask_dataset, train=True
+            )
+            self.normalize_label(multitask_dataset)
+
+        return multitask_dataset
+
     def _ready_to_load_all_from_file(self) -> bool:
         """
         Check if the data for all stages is ready to be loaded from files
@@ -1240,15 +1265,9 @@ def _save_data_to_files(self) -> None:
             stage: self._make_multitask_dataset(stage, save_smiles_and_ids=False, load_from_file=False)
             for stage in stages
         }
-
-        self.get_label_statistics(
-            self.processed_graph_data_path, self.data_hash, temp_datasets["train"], train=True
-        )
-        self.normalize_label(temp_datasets["train"])
         for stage in stages:
             self.save_featurized_data(temp_datasets[stage], self._path_to_load_from_file(stage))
             temp_datasets[stage].save_metadata(self._path_to_load_from_file(stage))
-
         # self.train_ds, self.val_ds, self.test_ds will be created during `setup()`
         del temp_datasets
 
@@ -1256,30 +1275,80 @@ 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_label_statistics(self, data_path, data_hash, dataset, train=False):
-        path_with_hash = os.path.join(data_path, data_hash)
-        os.makedirs(path_with_hash, exist_ok=True)
-        filename = os.path.join(path_with_hash, "task_norms.pkl")
-        # if self.task_norms was obtained from prepare_data step and
-        # this is training dataset split and the hash specific task_norms.pkl
-        # file does not exist, we recalculate the label statistics.
-        if self.task_norms and train and not os.path.isfile(filename):
-            for task in dataset[0]["labels"].keys():
-                labels = np.stack(np.array([datum["labels"][task] for datum in dataset]), axis=0)
+    def calculate_statistics(self, dataset: Datasets.MultitaskDataset, train: bool = False):
+        """
+        Calculate the statistics of the labels for each task, and overwrites the `self.task_norms` attribute.
+
+        Parameters:
+            dataset: the dataset to calculate the statistics from
+            train: whether the dataset is the training set
+
+        """
+        if self.task_norms and train:
+            for task in dataset.labels_size.keys():
+                # if the label type is graph_*, we need to stack them as the tensor shape is (num_labels, )
+                if task.startswith("graph"):
+                    labels = np.stack(
+                        np.array([datum["labels"][task] for datum in dataset if task in datum["labels"]]),
+                        axis=0,
+                    )
+                # for other tasks with node_ and edge_, the label shape is [num_nodes/num_edges, num_labels]
+                # we can concatenate them directly
+                else:
+                    labels = np.concatenate(
+                        [datum["labels"][task] for datum in dataset if task in datum["labels"]], axis=0
+                    )
+
                 self.task_norms[task].calculate_statistics(labels)
-            torch.save(self.task_norms, filename, pickle_protocol=4)
-        # if any of the above three condition does not satisfy, we load from file.
+
+    def get_label_statistics(
+        self,
+        data_path: Union[str, os.PathLike],
+        data_hash: str,
+        dataset: Datasets.MultitaskDataset,
+        train: bool = False,
+    ):
+        """
+        Get the label statistics from the dataset, and save them to file, if needed.
+        `self.task_norms` will be modified in-place with the label statistics.
+
+        Parameters:
+            data_path: the path to save and load the label statistics to. If None, no saving and loading will be done.
+            data_hash: the hash of the dataset generated by `get_data_hash()`
+            dataset: the dataset to calculate the statistics from
+            train: whether the dataset is the training set
+
+        """
+        if data_path is None:
+            self.calculate_statistics(dataset, train=train)
         else:
-            self.task_norms = torch.load(filename)
+            path_with_hash = os.path.join(data_path, data_hash)
+            os.makedirs(path_with_hash, exist_ok=True)
+            filename = os.path.join(path_with_hash, "task_norms.pkl")
+            if self.task_norms and train and not os.path.isfile(filename):
+                self.calculate_statistics(dataset, train=train)
+                torch.save(self.task_norms, filename, pickle_protocol=4)
+            # if any of the above three condition does not satisfy, we load from file.
+            else:
+                self.task_norms = torch.load(filename)
+
+    def normalize_label(self, dataset: Datasets.MultitaskDataset) -> Datasets.MultitaskDataset:
+        """
+        Normalize the labels in the dataset using the statistics in `self.task_norms`.
+
+        Parameters:
+            dataset: the dataset to normalize the labels from
 
-    def normalize_label(self, dataset):
-        for task in dataset[0]["labels"].keys():
+        Returns:
+            the dataset with normalized labels
+        """
+        for task in dataset.labels_size.keys():
             for i in range(len(dataset)):
-                normalized_label = self.task_norms[task].normalize(dataset[i]["labels"][task])
-                dataset[i]["labels"][task] = normalized_label
+                if task in dataset[i]["labels"]:
+                    dataset[i]["labels"][task] = self.task_norms[task].normalize(dataset[i]["labels"][task])
         return dataset
 
-    def save_featurized_data(self, dataset, processed_data_path):
+    def save_featurized_data(self, dataset: Datasets.MultitaskDataset, processed_data_path):
         os.makedirs(processed_data_path)  # In case the len(dataset) is 0
         for i in range(0, len(dataset), 1000):
             os.makedirs(os.path.join(processed_data_path, format(i // 1000, "04d")), exist_ok=True)
@@ -1297,7 +1366,9 @@ def process_func(self, param):
         index, datum, folder = param
         filename = os.path.join(folder, format(index // 1000, "04d"), format(index, "07d") + ".pkl")
         torch.save(
-            {"graph_with_features": datum["features"], "labels": datum["labels"]}, filename, pickle_protocol=4
+            {"graph_with_features": datum["features"], "labels": datum["labels"]},
+            filename,
+            pickle_protocol=4,
         )
         return
 
@@ -1603,6 +1674,7 @@ def _load_from_cache(self):
     def _extract_smiles_labels(
         self,
         df: pd.DataFrame,
+        task_level: str,
         smiles_col: str = None,
         label_cols: List[str] = [],
         idx_col: str = None,
@@ -1643,10 +1715,18 @@ def _extract_smiles_labels(
         label_cols = check_arg_iterator(label_cols, enforce_type=list)
         smiles = df[smiles_col].values
         if len(label_cols) > 0:
-            labels = [pd.to_numeric(df[col], errors="coerce") for col in label_cols]
-            labels = np.stack(labels, axis=1)
+            if task_level == "graph":
+                labels = extract_labels(df, "graph", label_cols)
+            elif task_level == "node":
+                labels = extract_labels(df, "node", label_cols)
+            elif task_level == "edge":
+                labels = extract_labels(df, "edge", label_cols)
+            elif task_level == "nodepair":
+                labels = extract_labels(df, "nodepair", label_cols)
+            else:
+                raise ValueError(f"Unknown task level: {task_level}")
         else:
-            labels = np.zeros([len(smiles), 0])
+            labels = float("nan") + np.zeros([len(smiles), 0])
 
         indices = None  # What are indices for?
         if idx_col is not None:
@@ -1748,9 +1828,9 @@ def _get_split_indices(
                     f"file type `{file_type}` for `{splits_path}` not recognised, please use .pt, .csv or .tsv"
                 )
             train, val, test = split_names
-            train_indices = splits[train].astype("int").tolist()
-            val_indices = splits[val].astype("int").tolist()
-            test_indices = splits[test].astype("int").tolist()
+            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()
 
         # Filter train, val and test indices
         _, train_idx, _ = np.intersect1d(sample_idx, train_indices, return_indices=True)
@@ -2064,6 +2144,7 @@ def __init__(
                 "smiles_col": smiles_col,
                 "label_cols": label_cols,
                 "splits_path": splits_path,
+                "task_level": task_args["task_level"],
             }
             self.metadata[task_name] = this_metadata
 
@@ -2329,6 +2410,7 @@ def prepare_data(self):
             args = self.task_dataset_processing_params[task]
             smiles, labels, sample_idx, extras = self._extract_smiles_labels(
                 df,
+                task_level=args.task_level,
                 smiles_col=args.smiles_col,
                 label_cols=args.label_cols,
                 idx_col=args.idx_col,
@@ -2433,6 +2515,7 @@ def get_fake_graph(self):
 
         smiles, labels, sample_idx, extras = self._extract_smiles_labels(
             df,
+            task_level=args.task_level,
             smiles_col=args.smiles_col,
             label_cols=label_cols,
             idx_col=args.idx_col,
diff --git a/goli/data/dataset.py b/goli/data/dataset.py
index 30d50212c..a648b95de 100644
--- a/goli/data/dataset.py
+++ b/goli/data/dataset.py
@@ -19,7 +19,7 @@
 class SingleTaskDataset(Dataset):
     def __init__(
         self,
-        labels: Union[torch.Tensor, np.ndarray],
+        labels: List[Union[torch.Tensor, np.ndarray]],
         features: Optional[List[Union[Data, "GraphDict"]]] = None,
         smiles: Optional[List[str]] = None,
         indices: Optional[List[str]] = None,
@@ -200,8 +200,6 @@ def __init__(
             if not save_smiles_and_ids:
                 self.mol_ids = None
                 self.smiles = None
-
-            self.labels = np.array(self.labels)
             self.labels_size = self.set_label_size_dict(datasets)
             self.dataset_length = len(self.labels)
             self._num_nodes_list = None
@@ -222,8 +220,8 @@ def save_metadata(self, directory: str):
             "smiles",
             "labels_size",
             "dataset_length",
-            "num_nodes_list",
-            "num_edges_list",
+            "_num_nodes_list",
+            "_num_edges_list",
         ]
         attrs = {attr: getattr(self, attr) for attr in attrs_to_save}
 
@@ -240,8 +238,8 @@ def _load_metadata(self):
             "smiles",
             "labels_size",
             "dataset_length",
-            "num_nodes_list",
-            "num_edges_list",
+            "_num_nodes_list",
+            "_num_edges_list",
         ]
         path = os.path.join(self.data_path, "multitask_metadata.pkl")
         attrs = torch.load(path)
@@ -369,6 +367,13 @@ def __getitem__(self, idx):
         if self.load_from_file:
             data_dict = self.load_graph_from_index(idx)
             datum["features"] = data_dict["graph_with_features"]
+            # these type changes are temporary for ipu dtype match
+            if hasattr(data_dict["labels"], "graph_pcba_1328"):
+                data_dict["labels"].graph_pcba_1328 = data_dict["labels"].graph_pcba_1328.astype(np.float32)
+            if hasattr(data_dict["labels"], "graph_pcqm4m_g25"):
+                data_dict["labels"].graph_pcqm4m_g25 = data_dict["labels"].graph_pcqm4m_g25.astype(np.float32)
+            if hasattr(data_dict["labels"], "node_pcqm4m_n4"):
+                data_dict["labels"].node_pcqm4m_n4 = data_dict["labels"].node_pcqm4m_n4.astype(np.float32)
             datum["labels"] = data_dict["labels"]
             if "smiles" in data_dict.keys():
                 datum["smiles"] = data_dict["smiles"]
@@ -431,12 +436,19 @@ def merge(
             smiles[unique_idx].append(all_lists["smiles"][all_idx])
 
         # Store the labels.
-        labels = [{} for _ in range(len(mol_ids))]
+        labels = [Data() for _ in range(len(mol_ids))]
         for all_idx, unique_idx in enumerate(inv):
-            task = all_lists["tasks"][all_idx]
+            task: str = all_lists["tasks"][all_idx]
             label = all_lists["labels"][all_idx]
             labels[unique_idx][task] = label
 
+            if all_idx < len(all_lists["features"]):
+                features = all_lists["features"][all_idx]
+                labels[unique_idx]["x"] = torch.empty(
+                    (features.num_nodes, 1)
+                )  # IPU is not happy with zero-sized tensors, so use shape (features.num_nodes, 1) here
+                labels[unique_idx]["edge_index"] = torch.empty((2, features.num_edges))
+
         # Store the features
         if len(all_lists["features"]) > 0:
             features = [-1 for i in range(len(mol_ids))]
@@ -474,7 +486,6 @@ def _get_all_lists_ids(self, datasets: Dict[str, SingleTaskDataset]) -> Dict[str
                 ds_features = [ds[i]["features"] for i in range(len(ds))]
             else:
                 ds_features = None
-
             all_smiles.extend(ds_smiles)
             all_labels.extend(ds_labels)
             all_mol_ids.extend(ds_mol_ids)
@@ -506,11 +517,20 @@ def set_label_size_dict(self, datasets: Dict[str, SingleTaskDataset]):
         for task, ds in datasets.items():
             if len(ds) == 0:
                 continue
-            label = ds[0][
-                "labels"
-            ]  # Assume for a fixed task, the label dimension is the same across data points, so we can choose the first data point for simplicity.
-            torch_label = torch.as_tensor(label)
-            # torch_label = label
+
+            valid_label = None
+            for i in range(len(ds)):
+                if ds[i] is not None:
+                    valid_label = ds[i]["labels"]
+                    break
+
+            if valid_label is None:
+                raise ValueError(f"Dataset for task {task} has no valid labels.")
+
+            # Assume for a fixed task, the label dimension is the same across data points
+            torch_label = torch.as_tensor(valid_label)
+
+            # First dimension is graph-specific
             task_labels_size[task] = torch_label.size()
         return task_labels_size
 
@@ -594,7 +614,7 @@ def __init__(
                 self.mol_ids, self.smiles, self.labels, _ = self.deepcopy_mol(
                     self.mol_ids, self.smiles, self.labels
                 )
-        self.labels = np.array(self.labels)
+
         self.labels_size = self.set_label_size_dict(datasets)
         self.features = self.features
 
diff --git a/goli/data/multilevel_utils.py b/goli/data/multilevel_utils.py
new file mode 100644
index 000000000..de8ef96c1
--- /dev/null
+++ b/goli/data/multilevel_utils.py
@@ -0,0 +1,46 @@
+import pandas as pd
+import ast
+import numpy as np
+from typing import List
+import itertools
+import math
+
+
+def extract_labels(df: pd.DataFrame, task_level: str, label_cols: List[str]):
+    """Extracts labels in label_cols from dataframe df for a given task_level.
+    Returns a list of numpy arrays converted to the correct shape. Multiple
+    targets are concatenated for each graph.
+    """
+
+    def unpack(graph_data):
+        graph_data = pd.to_numeric(graph_data, errors="coerce")
+        if isinstance(graph_data, str):
+            graph_data_list = ast.literal_eval(graph_data)
+            return np.array(graph_data_list)
+        elif isinstance(graph_data, (int, float)):
+            return np.array([graph_data])
+        elif isinstance(graph_data, list):
+            return np.array(graph_data)
+        elif isinstance(graph_data, np.ndarray):
+            return graph_data
+        else:
+            raise ValueError(
+                f"Graph data should be one of str, float, int, list, np.ndarray, got {type(graph_data)}"
+            )
+
+    def unpack_column(data: pd.Series):
+        return data.apply(unpack)
+
+    def merge_columns(data: pd.Series):
+        data = data.to_list()
+        data = [np.array([np.nan]) if not isinstance(d, np.ndarray) and math.isnan(d) else d for d in data]
+        padded_data = itertools.zip_longest(*data, fillvalue=np.nan)
+        data = np.stack(padded_data, 1).T
+        return data
+
+    unpacked_df: pd.DataFrame = df[label_cols].apply(unpack_column)
+    output = unpacked_df.apply(merge_columns, axis="columns").to_list()
+
+    if task_level == "graph":
+        return np.concatenate(output)
+    return output
diff --git a/goli/data/normalization.py b/goli/data/normalization.py
index 2dcb4d778..dfe5006a4 100644
--- a/goli/data/normalization.py
+++ b/goli/data/normalization.py
@@ -45,10 +45,11 @@ def calculate_statistics(self, array):
         """
         Saves the normalization parameters (e.g. mean and variance) to the object.
         """
-        self.data_max = np.nanmax(array).tolist()
-        self.data_min = np.nanmin(array).tolist()
-        self.data_mean = np.nanmean(array).tolist()  # 5.380503871833475 for pcqm4mv2
-        self.data_std = np.nanstd(array).tolist()  # 1.17850688410978995 for pcqm4mv2
+        # add axis = 0 to make sure that the statistics for multiple column labels is a vector or list instead of a scalar
+        self.data_max = np.nanmax(array, axis=0).tolist()
+        self.data_min = np.nanmin(array, axis=0).tolist()
+        self.data_mean = np.nanmean(array, axis=0).tolist()  # 5.380503871833475 for pcqm4mv2
+        self.data_std = np.nanstd(array, axis=0).tolist()  # 1.17850688410978995 for pcqm4mv2
         if self.verbose:
             logger.info(f"Max value for normalization '{self.data_max}'")
             logger.info(f"Min value for normalization '{self.data_min}'")
@@ -88,8 +89,14 @@ def denormalize(self, input):
         if self.method is None:
             return input
         elif self.method == "normal":
-            return (input * self.data_std) + self.data_mean
+            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)
+            return (input * std) + mean
         elif self.method == "unit":
-            return input * (self.data_max - self.data_min) + self.data_min
+            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)
+            return input * (dmax - dmin) + dmin
         else:
             raise ValueError(f"normalization method {self.method} not recognised.")
diff --git a/goli/features/featurizer.py b/goli/features/featurizer.py
index 97785535c..eb9e1d12f 100644
--- a/goli/features/featurizer.py
+++ b/goli/features/featurizer.py
@@ -823,11 +823,15 @@ def mol_to_adjacency_matrix(
         adj_val.extend([val, val])
 
     # Convert to torch coo sparse tensor
-    adj = coo_matrix(
-        (torch.as_tensor(adj_val), torch.as_tensor(adj_idx).T.reshape(2, -1)),
-        shape=(mol.GetNumAtoms(), mol.GetNumAtoms()),
-        dtype=dtype,
-    )
+    if len(adj_val) > 0:  # ensure tensor is not empty
+        adj = coo_matrix(
+            (torch.as_tensor(adj_val), torch.as_tensor(adj_idx).T.reshape(2, -1)),
+            shape=(mol.GetNumAtoms(), mol.GetNumAtoms()),
+            dtype=dtype,
+        )
+    else:
+        # Special case for molecules with one atom
+        adj = coo_matrix(([], np.array([[], []])), shape=(mol.GetNumAtoms(), mol.GetNumAtoms()), dtype=dtype)
 
     # Add self loops
     if add_self_loop:
diff --git a/goli/ipu/ipu_dataloader.py b/goli/ipu/ipu_dataloader.py
index 76a439b33..1b9510c75 100644
--- a/goli/ipu/ipu_dataloader.py
+++ b/goli/ipu/ipu_dataloader.py
@@ -149,6 +149,7 @@ def __call__(
             )
 
             local_batch["features"] = transform(local_batch["features"])
+            local_batch["labels"] = transform(local_batch["labels"])
             all_batches.append(local_batch)
 
         out_batch = {}
@@ -156,7 +157,7 @@ def __call__(
         # Stack tensors in the first dimension to allow IPUs to differentiate between local and global graph
         out_batch["labels"] = {
             key: torch.stack([this_batch["labels"][key] for this_batch in all_batches], 0)
-            for key in all_batches[0]["labels"].keys()
+            for key in all_batches[0]["labels"].keys
         }
         out_graphs = [this_batch["features"] for this_batch in all_batches]
         stacked_features = deepcopy(out_graphs[0])
@@ -342,6 +343,12 @@ def validate(self, data):
         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.
@@ -387,6 +394,11 @@ def __call__(self, batch: Batch) -> Batch:
                     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
 
diff --git a/goli/ipu/ipu_losses.py b/goli/ipu/ipu_losses.py
index 2ba2f39c3..6cb6fa907 100644
--- a/goli/ipu/ipu_losses.py
+++ b/goli/ipu/ipu_losses.py
@@ -2,6 +2,8 @@
 from torch import Tensor
 from torch.nn import BCELoss, MSELoss, L1Loss
 from torch._C import _infer_size
+from loguru import logger
+from goli.trainer.losses import HybridCELoss
 
 
 class BCELossIPU(BCELoss):
@@ -14,7 +16,7 @@ class BCELossIPU(BCELoss):
     def forward(self, input: Tensor, target: Tensor) -> Tensor:
         prev_weight = None
 
-        target = target.clone()
+        target = target.clone().to(input.dtype)
         weight = self.weight
 
         # Get the original weight matrix. If None, set all weights = 1
@@ -37,10 +39,15 @@ def forward(self, input: Tensor, target: Tensor) -> Tensor:
         # Compute the loss, and rescale by the number of nan elements
         self.weight = weight
         loss = super().forward(input, target)
-        loss = loss * nan_targets.numel() / ((~nan_targets).sum())
+
+        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
 
 
@@ -53,7 +60,7 @@ class MSELossIPU(MSELoss):
     """
 
     def forward(self, input: Tensor, target: Tensor) -> Tensor:
-        target = target.clone()
+        target = target.clone().to(input.dtype)
         input = input.clone()
 
         # Replace the nan-targets in the input/target tensors by 0
@@ -63,7 +70,11 @@ def forward(self, input: Tensor, target: Tensor) -> Tensor:
 
         # Compute the loss, and rescale by the number of nan elements
         loss = super().forward(input, target)
-        loss = loss * nan_targets.numel() / ((~nan_targets).sum())
+
+        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
 
@@ -77,7 +88,33 @@ class L1LossIPU(L1Loss):
     """
 
     def forward(self, input: Tensor, target: Tensor) -> Tensor:
-        target = target.clone()
+        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 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
@@ -87,6 +124,9 @@ def forward(self, input: Tensor, target: Tensor) -> Tensor:
 
         # Compute the loss, and rescale by the number of nan elements
         loss = super().forward(input, target)
-        loss = loss * nan_targets.numel() / ((~nan_targets).sum())
+        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
diff --git a/goli/ipu/ipu_metrics.py b/goli/ipu/ipu_metrics.py
index 3f8b10dc2..197afd0b4 100644
--- a/goli/ipu/ipu_metrics.py
+++ b/goli/ipu/ipu_metrics.py
@@ -1,4 +1,4 @@
-from typing import Optional, Tuple, Sequence
+from typing import Optional, Tuple, Sequence, Literal
 
 import torch
 from torch import BoolTensor, IntTensor, Tensor
@@ -24,6 +24,7 @@ def auroc_ipu(
     preds: Tensor,
     target: Tensor,
     num_classes: Optional[int] = None,
+    task: Optional[Literal["binary", "multiclass", "multilabel"]] = None,
     pos_label: Optional[int] = None,
     average: Optional[str] = "macro",
     max_fpr: Optional[float] = None,
@@ -54,6 +55,7 @@ def auroc_ipu(
         preds=preds,
         target=target.to(int),
         num_classes=num_classes,
+        task=task,
         pos_label=pos_label,
         average=average,
         max_fpr=max_fpr,
@@ -67,6 +69,8 @@ def average_precision_ipu(
     preds: Tensor,
     target: Tensor,
     num_classes: Optional[int] = None,
+    task: Optional[str] = "multiclass",
+    ignore_index: Optional[int] = None,
     pos_label: Optional[int] = None,
     average: Optional[str] = "macro",
     sample_weights: Optional[Sequence] = None,
@@ -98,6 +102,8 @@ def average_precision_ipu(
         preds=preds,
         target=target,
         num_classes=num_classes,
+        task=task,
+        ignore_index=ignore_index,
         pos_label=pos_label,
         average=average,
         # sample_weights=sample_weights,
@@ -559,7 +565,7 @@ def pearson_ipu(preds, target):
     preds = NaNTensor(preds)
     target = NaNTensor(target)
     preds[target.get_nans] = float("nan")
-    pearson = pearson_corrcoef(preds, target)
+    pearson = pearson_corrcoef(preds, target.to(preds.dtype))
     return Tensor(pearson)
 
 
diff --git a/goli/ipu/ipu_utils.py b/goli/ipu/ipu_utils.py
index d5d732329..11af049b3 100644
--- a/goli/ipu/ipu_utils.py
+++ b/goli/ipu/ipu_utils.py
@@ -1,3 +1,4 @@
+import os
 import tempfile
 from datetime import datetime
 from copy import deepcopy
@@ -101,6 +102,7 @@ def load_ipu_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)
@@ -142,7 +144,8 @@ def ipu_options_list_to_file(ipu_opts: Optional[List[str]]) -> tempfile._Tempora
     if ipu_opts is None:
         return
 
-    tmp_file = tempfile.NamedTemporaryFile("w")
+    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/goli/ipu/ipu_wrapper.py b/goli/ipu/ipu_wrapper.py
index 3c56e91ae..3709e526d 100644
--- a/goli/ipu/ipu_wrapper.py
+++ b/goli/ipu/ipu_wrapper.py
@@ -19,19 +19,6 @@
 poptorch = import_poptorch()
 
 
-def remove_pad_loss(preds: Dict[str, Tensor], targets: Dict[str, Tensor]):
-    """
-    helper function to remove the fake graph loss
-    always reduce the last loss since it is the fake graph
-    """
-    for task in targets.keys():
-        if targets[task].shape == preds[task].shape:
-            continue
-        else:
-            preds[task] = preds[task][:-1]
-    return preds
-
-
 class DictIPUStrategy(IPUStrategy):
     def _step(self, stage: RunningStage, *args: Any, **kwargs: Any) -> STEP_OUTPUT:
         args = self._prepare_input(args)
@@ -115,11 +102,9 @@ def compute_loss(
         targets: Dict[str, Tensor],
         weights: Optional[Tensor],
         loss_fun: Dict[str, Callable],
-        target_nan_mask: Union[Type, str] = "ignore-flatten",
+        target_nan_mask: Union[Type, str] = "ignore",
         multitask_handling: Optional[str] = None,
     ) -> Tuple[Tensor, Dict[str, Tensor]]:
-        preds = remove_pad_loss(preds, targets)
-
         return PredictorModule.compute_loss(
             preds, targets, weights, loss_fun, target_nan_mask, multitask_handling
         )
diff --git a/goli/nn/base_layers.py b/goli/nn/base_layers.py
index d60dc0a28..d0206a93a 100644
--- a/goli/nn/base_layers.py
+++ b/goli/nn/base_layers.py
@@ -149,103 +149,49 @@ def forward(
         **kwargs,
     ) -> Tuple[Tensor, Optional[Tensor]]:
         # attn_bias [batch, num_heads, nodes, nodes]
-        if self.biased_attention and attn_bias is not None:
-            # assuming source and target have the same sequence length (homogeneous graph attention)
-            batch, nodes, hidden = query.size()
-            assert (
-                hidden == self.embed_dim
-            ), f"query hidden dimension {hidden} != embed_dim {self.embed_dim} in class"
-            head_dim = self.embed_dim // self.num_heads
-            assert head_dim * self.num_heads == self.embed_dim, "embed_dim must be divisible by num_heads"
-            scaling_factor = 1 / head_dim  # use head_dim instead of (head_dim**0.5) for mup
-            b_q, b_k, b_v = self.in_proj_bias.chunk(3)
-            q_proj_weight, k_proj_weight, v_proj_weight = self.in_proj_weight.chunk(3)
-            # [batch, num_heads, nodes, head_size]
-            q = linear(query, q_proj_weight, b_q).view(batch, nodes, self.num_heads, -1).transpose(1, 2)
-            # [batch, num_heads, nodes, head_size]
-            k = linear(key, k_proj_weight, b_k).view(batch, nodes, self.num_heads, -1).transpose(1, 2)
-            # [batch, num_heads, nodes, head_size]
-            v = linear(value, v_proj_weight, b_v).view(batch, nodes, self.num_heads, -1).transpose(1, 2)
-            q = q * scaling_factor
-            # [batch, num_heads, nodes, nodes]
-            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),
-                    key_padding_mask_value,
-                )
-            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]
-            attn = attn_probs @ v
-            # [batch, nodes, embd_dim]
-            attn = attn.transpose(1, 2).contiguous().view(batch, nodes, self.embed_dim)
-            # [batch, nodes, embd_dim]
-            out = (self.out_proj(attn), None)
-        else:
-            # Patching the forward to use a different scaling for the dot-product
-            prev_fn = F._scaled_dot_product_attention
-            F._scaled_dot_product_attention = _mup_scaled_dot_product_attention
-            out = super().forward(
-                query=query, key=key, value=value, key_padding_mask=key_padding_mask, *args, **kwargs
+        if not self.biased_attention or attn_bias is None:
+            attn_bias = 0.0
+        # assuming source and target have the same sequence length (homogeneous graph attention)
+        batch, nodes, hidden = query.size()
+        assert (
+            hidden == self.embed_dim
+        ), f"query hidden dimension {hidden} != embed_dim {self.embed_dim} in class"
+        head_dim = self.embed_dim // self.num_heads
+        assert head_dim * self.num_heads == self.embed_dim, "embed_dim must be divisible by num_heads"
+        scaling_factor = 1 / head_dim  # use head_dim instead of (head_dim**0.5) for mup
+        b_q, b_k, b_v = self.in_proj_bias.chunk(3)
+        q_proj_weight, k_proj_weight, v_proj_weight = self.in_proj_weight.chunk(3)
+        # [batch, num_heads, nodes, head_size]
+        q = linear(query, q_proj_weight, b_q).view(batch, nodes, self.num_heads, -1).transpose(1, 2)
+        # [batch, num_heads, nodes, head_size]
+        k = linear(key, k_proj_weight, b_k).view(batch, nodes, self.num_heads, -1).transpose(1, 2)
+        # [batch, num_heads, nodes, head_size]
+        v = linear(value, v_proj_weight, b_v).view(batch, nodes, self.num_heads, -1).transpose(1, 2)
+        q = q * scaling_factor
+        # [batch, num_heads, nodes, nodes]
+        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),
+                key_padding_mask_value,
             )
-            F._scaled_dot_product_attention = prev_fn
+        else:
+            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]
+        attn = attn_probs @ v
+        # [batch, nodes, embd_dim]
+        attn = attn.transpose(1, 2).contiguous().view(batch, nodes, self.embed_dim)
+        # [batch, nodes, embd_dim]
+        out = (self.out_proj(attn), None)
         return out
 
 
-def _mup_scaled_dot_product_attention(
-    q: Tensor,
-    k: Tensor,
-    v: Tensor,
-    attn_mask: Optional[Tensor] = None,
-    key_padding_mask: Optional[Tensor] = None,
-    dropout_p: float = 0.0,
-) -> Tuple[Tensor, Tensor]:
-    r"""
-    Computes scaled dot product attention on query, key and value tensors, using
-    an optional attention mask if passed, and applying dropout if a probability
-    greater than 0.0 is specified.
-    Returns a tensor pair containing attended values and attention weights.
-
-    This modifies the standard torch function by normalizing with `1/d` instead of `1/sqrt(d)`
-
-    Args:
-        q, k, v: query, key and value tensors. See Shape section for shape details.
-        attn_mask: optional tensor containing mask values to be added to calculated
-            attention. May be 2D or 3D; see Shape section for details.
-        dropout_p: dropout probability. If greater than 0.0, dropout is applied.
-
-    Shape:
-        - q: :math:`(B, Nt, E)` where B is batch size, Nt is the target sequence length,
-            and E is embedding dimension.
-        - key: :math:`(B, Ns, E)` where B is batch size, Ns is the source sequence length,
-            and E is embedding dimension.
-        - value: :math:`(B, Ns, E)` where B is batch size, Ns is the source sequence length,
-            and E is embedding dimension.
-        - attn_mask: either a 3D tensor of shape :math:`(B, Nt, Ns)` or a 2D tensor of
-            shape :math:`(Nt, Ns)`.
-
-        - Output: attention values have shape :math:`(B, Nt, E)`; attention weights
-            have shape :math:`(B, Nt, Ns)`
-    """
-    B, Nt, E = q.shape
-    q = q / E  # Instead of `q / math.sqrt(E)`
-    # (B, Nt, E) x (B, E, Ns) -> (B, Nt, Ns)
-    attn = torch.bmm(q, k.transpose(-2, -1))
-    if attn_mask is not None:
-        attn += attn_mask
-    attn = F.softmax(attn, dim=-1)
-    if dropout_p > 0.0:
-        attn = F.dropout(attn, p=dropout_p)
-    # (B, Nt, Ns) x (B, Ns, E) -> (B, Nt, E)
-    output = torch.bmm(attn, v)
-    return output, attn
-
-
 class TransformerEncoderLayerMup(nn.TransformerEncoderLayer):
     r"""
     Modified version of ``torch.nn.TransformerEncoderLayer`` that uses :math:`1/n`-scaled attention
diff --git a/goli/nn/pyg_layers/__init__.py b/goli/nn/pyg_layers/__init__.py
index 70f90d8c4..d5ffb083d 100644
--- a/goli/nn/pyg_layers/__init__.py
+++ b/goli/nn/pyg_layers/__init__.py
@@ -1,3 +1,4 @@
+from .gcn_pyg import GCNConvPyg
 from .gated_gcn_pyg import GatedGCNPyg
 from .gin_pyg import GINConvPyg
 from .gin_pyg import GINEConvPyg
diff --git a/goli/nn/pyg_layers/gcn_pyg.py b/goli/nn/pyg_layers/gcn_pyg.py
new file mode 100644
index 000000000..310aa8f7a
--- /dev/null
+++ b/goli/nn/pyg_layers/gcn_pyg.py
@@ -0,0 +1,111 @@
+from typing import Callable, Union
+from functools import partial
+
+import torch_geometric.nn as pyg_nn
+from torch_geometric.data import Data, Batch
+
+from goli.nn.base_graph_layer import BaseGraphModule, check_intpus_allow_int
+from goli.utils.decorators import classproperty
+
+
+class GCNConvPyg(BaseGraphModule):
+    def __init__(
+        self,
+        in_dim: int,
+        out_dim: int,
+        activation: Union[Callable, str] = "relu",
+        dropout: float = 0.0,
+        normalization: Union[str, Callable] = "none",
+        **kwargs,
+    ):
+        super().__init__(
+            in_dim=in_dim,
+            out_dim=out_dim,
+            activation=activation,
+            dropout=dropout,
+            normalization=normalization,
+            **kwargs,
+        )
+
+        self.model = pyg_nn.GCNConv(
+            in_channels=self.in_dim, out_channels=out_dim, add_self_loops=False, normalize=False
+        )
+        self.model.__check_input__ = partial(check_intpus_allow_int, self)
+
+    def forward(
+        self,
+        batch: Union[Data, Batch],
+    ) -> Union[Data, Batch]:
+        r"""
+        forward function of the layer
+        Parameters:
+            batch: pyg Batch graphs to pass through the layer
+        Returns:
+            batch: pyg Batch graphs
+        """
+        batch.feat = self.model(batch.feat, batch.edge_index)
+        batch.feat = self.apply_norm_activation_dropout(batch.feat, batch_idx=batch.batch)
+
+        return batch
+
+    @classproperty
+    def layer_supports_edges(cls) -> bool:
+        r"""
+        Return a boolean specifying if the layer type supports edges or not.
+
+        Returns:
+
+            supports_edges: bool
+                Always ``False`` for the current class
+        """
+        return False
+
+    @property
+    def layer_inputs_edges(self) -> bool:
+        r"""
+        Return a boolean specifying if the layer type
+        uses edges as input or not.
+        It is different from ``layer_supports_edges`` since a layer that
+        supports edges can decide to not use them.
+
+        Returns:
+
+            bool:
+                Always ``False`` for the current class
+        """
+        return False
+
+    @property
+    def layer_outputs_edges(self) -> bool:
+        r"""
+        Abstract method. Return a boolean specifying if the layer type
+        uses edges as input or not.
+        It is different from ``layer_supports_edges`` since a layer that
+        supports edges can decide to not use them.
+
+        Returns:
+
+            bool:
+                Always ``False`` for the current class
+        """
+        return False
+
+    @property
+    def out_dim_factor(self) -> int:
+        r"""
+        Get the factor by which the output dimension is multiplied for
+        the next layer.
+
+        For standard layers, this will return ``1``.
+
+        But for others, such as ``GatLayer``, the output is the concatenation
+        of the outputs from each head, so the out_dim gets multiplied by
+        the number of heads, and this function should return the number
+        of heads.
+
+        Returns:
+
+            int:
+                Always ``1`` for the current class
+        """
+        return 1
diff --git a/goli/trainer/losses.py b/goli/trainer/losses.py
index eddec6993..9392b0d8c 100644
--- a/goli/trainer/losses.py
+++ b/goli/trainer/losses.py
@@ -9,7 +9,7 @@
 class HybridCELoss(_WeightedLoss):
     def __init__(
         self,
-        n_brackets: int,
+        n_brackets: int = 5,
         regression_loss: str = "mse",
         alpha: float = 0.5,
         weight: Optional[Tensor] = None,
@@ -57,12 +57,10 @@ def forward(self, input: Tensor, target: Tensor) -> Tensor:
             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}.
         """
-        if self.brackets.device != input.device:
-            self.brackets = self.brackets.to(input.device)
 
         target = target.flatten()
 
-        regression_input = torch.inner(input, self.brackets)
+        regression_input = torch.inner(input, self.brackets.to(input.device))
         regression_loss = self.regression_loss(regression_input, target.float(), reduction=self.reduction)
 
         ce_loss = F.cross_entropy(input, target.long(), weight=self.weight, reduction=self.reduction)
diff --git a/goli/trainer/metrics.py b/goli/trainer/metrics.py
index f228dee35..6e3148194 100644
--- a/goli/trainer/metrics.py
+++ b/goli/trainer/metrics.py
@@ -265,11 +265,17 @@ def compute(self, preds: Tensor, target: Tensor) -> Tensor:
         if self.thresholder is not None:
             preds, target = self.thresholder(preds, target)
 
-        if self.target_to_int:
-            target = target.to(int)
-
         target_nans = torch.isnan(target)
 
+        # for the classifigression task, cast predictions from
+        # (batch_size, n_targets * n_brackets) to (batch_size, n_targets, n_brackets)
+        # TODO: make this more flexible to the target shape in the future
+        if preds.shape[1] != target.shape[1]:
+            preds = preds.view(target.shape[0], target.shape[1], -1)
+            classifigression = True
+        else:
+            classifigression = False
+
         if self.multitask_handling is None:
             # In case of no multi-task handling, apply the nan filtering, then compute the metrics
             assert (
@@ -278,24 +284,38 @@ def compute(self, preds: Tensor, target: Tensor) -> Tensor:
             preds, target = self._filter_nans(preds, target)
             if self.squeeze_targets:
                 target = target.squeeze()
+            if self.target_to_int:
+                target = target.to(int)
             metric_val = self.metric(preds, target, **self.kwargs)
         elif self.multitask_handling == "flatten":
             # Flatten the tensors, apply the nan filtering, then compute the metrics
-            preds, target = preds.flatten(), target.flatten()
+            if classifigression:
+                preds = preds.view(-1, preds.shape[-1])
+                target = target.flatten()
+            else:
+                preds, target = preds.flatten(), target.flatten()
             preds, target = self._filter_nans(preds, target)
             if self.squeeze_targets:
                 target = target.squeeze()
+            if self.target_to_int:
+                target = target.to(int)
             metric_val = self.metric(preds, target, **self.kwargs)
         elif self.multitask_handling == "mean-per-label":
             # Loop the columns (last dim) of the tensors, apply the nan filtering, compute the metrics per column, then average the metrics
             target_list = [target[..., ii][~target_nans[..., ii]] for ii in range(target.shape[-1])]
-            preds_list = [preds[..., ii][~target_nans[..., ii]] for ii in range(preds.shape[-1])]
+            # TODO: make this more flexible to the target shape in the future
+            if classifigression:
+                preds_list = [preds[..., i, :][~target_nans[..., i]] for i in range(preds.shape[1])]
+            else:
+                preds_list = [preds[..., ii][~target_nans[..., ii]] for ii in range(preds.shape[-1])]
             metric_val = []
             for ii in range(len(target_list)):
                 try:
                     this_preds, this_target = self._filter_nans(preds_list[ii], target_list[ii])
                     if self.squeeze_targets:
                         this_target = this_target.squeeze()
+                    if self.target_to_int:
+                        this_target = this_target.to(int)
                     metric_val.append(self.metric(this_preds, this_target, **self.kwargs))
                 except:
                     pass
diff --git a/goli/trainer/predictor.py b/goli/trainer/predictor.py
index eb6da6ecf..74fee8cee 100644
--- a/goli/trainer/predictor.py
+++ b/goli/trainer/predictor.py
@@ -57,6 +57,7 @@ def __init__(
             metrics_on_progress_bar: A `dict[str, list[str2]`, where `str` is the task name and `str2` the metrics to include on the progress bar
             metrics_on_training_set: A `dict[str, list[str2]`, where `str` is the task name and `str2` the metrics to include on the training set
             flag_kwargs: Arguments related to using the FLAG adversarial augmentation
+            task_norms: the normalization for each task
         """
         self.save_hyperparameters()
 
@@ -66,10 +67,12 @@ def __init__(
 
         self.target_nan_mask = target_nan_mask
         self.multitask_handling = multitask_handling
+        self.task_norms = task_norms
 
         super().__init__()
 
         # Setting the model options
+        self.model_kwargs = model_kwargs
         self._model_options = ModelOptions(model_class=model_class, model_kwargs=model_kwargs)
         # Setting the optimizer options
         self.optim_options = OptimOptions(
@@ -91,10 +94,22 @@ def __init__(
             eval_options[task].check_metrics_validity()
 
         self._eval_options_dict: Dict[str, EvalOptions] = eval_options
+        self._eval_options_dict = {
+            self._get_task_key(
+                task_level=model_kwargs["task_heads_kwargs"][key]["task_level"], task=key
+            ): value
+            for key, value in self._eval_options_dict.items()
+        }
         # Setting the flag options
         self._flag_options = FlagOptions(flag_kwargs=flag_kwargs)
 
         self.model = self._model_options.model_class(**self._model_options.model_kwargs)
+        loss_fun = {
+            self._get_task_key(
+                task_level=model_kwargs["task_heads_kwargs"][key]["task_level"], task=key
+            ): value
+            for key, value in loss_fun.items()
+        }
         self.tasks = list(loss_fun.keys())
 
         # Task-specific evalutation attributes
@@ -135,7 +150,6 @@ def __init__(
             task_metrics_on_progress_bar=self.metrics_on_progress_bar,
             monitor=monitor,
             mode=mode,
-            task_norms=task_norms,
         )
 
         # This helps avoid a bug when saving hparams to yaml with different dict or str formats
@@ -180,6 +194,12 @@ def _convert_features_dtype(self, feats):
                     feats[key] = val.to(dtype=self.dtype)
         return feats
 
+    def _get_task_key(self, task_level: str, task: str):
+        task_prefix = f"{task_level}_"
+        if not task.startswith(task_prefix):
+            task = task_prefix + task
+        return task
+
     def configure_optimizers(self, impl=None):
         if impl is None:
             impl = torch.optim.Adam
@@ -282,8 +302,20 @@ def _general_step(self, batch: Dict[str, Tensor], step_name: str, to_cpu: bool)
         elif isinstance(preds, Tensor):
             preds = {k: preds[ii] for ii, k in enumerate(targets_dict.keys())}
 
+        preds = {
+            self._get_task_key(
+                task_level=self.model_kwargs["task_heads_kwargs"][key]["task_level"], task=key
+            ): value
+            for key, value in preds.items()
+        }
         # preds = {k: preds[ii] for ii, k in enumerate(targets_dict.keys())}
         for task, pred in preds.items():
+            task_specific_norm = self.task_norms[task] if self.task_norms is not None else None
+            if step_name != "train":
+                # apply denormalization for val and test predictions for correct loss and metrics evaluation
+                # targets for val and test were not normalized
+                # train loss will stay as the normalized version
+                preds[task] = task_specific_norm.denormalize(pred)
             targets_dict[task] = targets_dict[task].to(dtype=pred.dtype)
         weights = batch.get("weights", None)
 
@@ -298,6 +330,12 @@ def _general_step(self, batch: Dict[str, Tensor], step_name: str, to_cpu: bool)
 
         device = "cpu" if to_cpu else None
         for task in preds:
+            task_specific_norm = self.task_norms[task] if self.task_norms is not None else None
+            if step_name == "train":
+                # apply denormalization for targets and predictions for the evaluation of metrics (excluding loss)
+                # train loss will stay as the normalized version
+                preds[task] = task_specific_norm.denormalize(preds[task])
+                targets_dict[task] = task_specific_norm.denormalize(targets_dict[task])
             preds[task] = preds[task].detach().to(device=device)
             targets_dict[task] = targets_dict[task].detach().to(device=device)
         if weights is not None:
diff --git a/goli/trainer/predictor_summaries.py b/goli/trainer/predictor_summaries.py
index 6894951a9..0c3fb735c 100644
--- a/goli/trainer/predictor_summaries.py
+++ b/goli/trainer/predictor_summaries.py
@@ -41,7 +41,6 @@ def __init__(
         monitor: str = "loss",
         mode: str = "min",
         task_name: Optional[str] = None,
-        task_specific_norm: Optional[Union[Callable, Any]] = None,
     ):
         r"""
         A container to be used by the Predictor Module that stores the results for the given metrics on the predictions and targets provided.
@@ -68,9 +67,6 @@ def __init__(
             task_name:
             name of the task (Default=`None`)
 
-            task_specific_norm:
-            task specific normalization (Default=`None`)
-
         """
         self.loss_fun = loss_fun
         self.metrics = metrics
@@ -91,7 +87,6 @@ def __init__(
         self.n_epochs: int = None
 
         self.task_name = task_name
-        self.task_specific_norm = task_specific_norm
         self.logged_metrics_exceptions = []  # Track which metric exceptions have been logged
 
     def update_predictor_state(
@@ -240,11 +235,6 @@ def get_metrics_logs(self) -> Dict[str, Any]:
             A dictionary of metrics to log.
         """
         targets = self.targets.to(dtype=self.predictions.dtype, device=self.predictions.device)
-        # apply denormalization for predictions
-        self.predictions = self.task_specific_norm.denormalize(self.predictions)
-        if self.step_name == "train":
-            # apply denormalization for targets
-            targets = self.task_specific_norm.denormalize(targets)
         # Compute the metrics always used in regression tasks
         metric_logs = {}
         metric_logs[self.metric_log_name(self.task_name, "mean_pred", self.step_name)] = nan_mean(
@@ -342,7 +332,6 @@ def __init__(
         task_metrics_on_progress_bar: List[str],
         monitor: str = "loss",
         mode: str = "min",
-        task_norms: Optional[Dict[Callable, Any]] = None,
     ):
         r"""
         class to store the summaries of the tasks
@@ -353,7 +342,6 @@ class to store the summaries of the tasks
             task_metrics_on_progress_bar: the metrics to use on the progress bar
             monitor: the metric to monitor
             mode: the mode of the metric to monitor
-            task_norms: the normalization for each task
         """
         self.task_loss_fun = task_loss_fun
         self.task_metrics = task_metrics
@@ -375,7 +363,6 @@ class to store the summaries of the tasks
                 self.monitor,
                 self.mode,
                 task_name=task,
-                task_specific_norm=task_norms[task] if task_norms is not None else None,
             )
 
         # Current predictor state
diff --git a/goli/utils/command_line_utils.py b/goli/utils/command_line_utils.py
new file mode 100644
index 000000000..ce02e3fb0
--- /dev/null
+++ b/goli/utils/command_line_utils.py
@@ -0,0 +1,73 @@
+import re
+from collections import defaultdict
+from typing import List, Dict
+
+
+def get_anchors_and_aliases(filepath):
+    """Utility function to extract anchors and aliases from YAML config file
+
+    In a YAML file we can specify an anchor as  `some_name: &anchor_name value`
+    This is then picked up with alises in the rest of the config as
+    `other_name: *anchor_name`
+    Using this format in the YAML file all anchors and aliases will be extracted
+
+    Args:
+        filepath (str): path to the config file
+
+    Returns:
+        anchors (Dict): A dictionary containing the YAML paths of the anchors
+                        as keys, and a list of YAML paths to alises.
+    """
+    anchors = defaultdict(list)
+    current_level = {}
+    anchor_to_path = {}
+
+    with open(filepath, "r") as file:
+        for line in file:
+            indent = len(line) - len(line.lstrip(" "))
+            key_match = re.search(r"(\w+):", line)
+            anchor_match = re.search(r"&(\w+)", line)
+            alias_match = re.search(r"\*(\w+)", line)
+            if key_match:
+                key = key_match.group(1)
+                # Compute the full path of the current key.
+                full_path = ".".join(
+                    [current_level[i] for i in sorted(current_level.keys()) if i < indent] + [key]
+                )
+                current_level[indent] = key
+                # Remove any keys that are indented more than the current line.
+                keys_to_remove = [i for i in current_level if i > indent]
+                for i in keys_to_remove:
+                    del current_level[i]
+            else:
+                full_path = ".".join([current_level[i] for i in sorted(current_level.keys())])
+            if anchor_match:
+                anchor = anchor_match.group(1)
+                anchor_to_path[anchor] = full_path
+            if alias_match:
+                alias = alias_match.group(1)
+                if alias in anchor_to_path:
+                    anchors[anchor_to_path[alias]].append(full_path)
+    return anchors
+
+
+def update_config(cfg: Dict, unknown: List, anchors: List):
+    """
+    Update the configuration dictionary with command line arguments.
+    """
+    for arg in unknown:
+        if arg.startswith("--"):
+            key, value = arg[2:].split("=")
+            if key in anchors.keys():
+                all_refs = anchors[key]
+            else:
+                all_refs = []
+            all_refs.append(key)
+            for key in all_refs:
+                keys = key.split(".")
+                temp_cfg = cfg
+                for k in keys[:-1]:
+                    temp_cfg = temp_cfg[k]
+                temp_cfg[keys[-1]] = type(temp_cfg[keys[-1]])(value) if keys[-1] in temp_cfg else value
+
+    return cfg
diff --git a/goli/utils/spaces.py b/goli/utils/spaces.py
index 1b4572d00..658434121 100644
--- a/goli/utils/spaces.py
+++ b/goli/utils/spaces.py
@@ -27,6 +27,7 @@
 }
 
 PYG_LAYERS_DICT = {
+    "pyg:gcn": PygLayers.GCNConvPyg,
     "pyg:gin": PygLayers.GINConvPyg,
     "pyg:gine": PygLayers.GINEConvPyg,
     "pyg:gated-gcn": PygLayers.GatedGCNPyg,
@@ -58,6 +59,7 @@
     "mse_ipu": IPULosses.MSELossIPU,
     "mae_ipu": IPULosses.L1LossIPU,
     "l1_ipu": IPULosses.L1LossIPU,
+    "hybrid_ce_ipu": IPULosses.HybridCELossIPU,
 }
 
 


Test patch:
diff --git a/expts/test_yaml.py b/expts/test_yaml.py
new file mode 100644
index 000000000..762364067
--- /dev/null
+++ b/expts/test_yaml.py
@@ -0,0 +1,79 @@
+# test_yaml.py
+
+
+import yaml
+
+CONFIG_FILE = "neurips2023_configs/config_small_mpnn.yaml"
+
+import re
+from collections import defaultdict
+
+# def get_anchors_and_aliases(filepath):
+#     anchors = defaultdict(list)
+#     current_level = {}
+
+#     with open(filepath, "r") as file:
+#         for line in file:
+#             indent = len(line) - len(line.lstrip(' '))
+#             key_match = re.search(r'(\w+):', line)
+#             anchor_match = re.search(r'&(\w+)', line)
+#             alias_match = re.search(r'\*(\w+)', line)
+#             if key_match:
+#                 key = key_match.group(1)
+#                 # Compute the full path of the current key.
+#                 full_path = '.'.join([current_level[i] for i in sorted(current_level.keys()) if i < indent] + [key])
+#                 current_level[indent] = key
+#                 # Remove any keys that are indented more than the current line.
+#                 keys_to_remove = [i for i in current_level if i > indent]
+#                 for i in keys_to_remove:
+#                     del current_level[i]
+#             else:
+#                 full_path = '.'.join([current_level[i] for i in sorted(current_level.keys())])
+#             if anchor_match:
+#                 anchor = anchor_match.group(1)
+#                 anchors[anchor].append(full_path)
+#             if alias_match:
+#                 alias = alias_match.group(1)
+#                 anchors[alias].append(full_path)
+#     return {k: v for k, v in anchors.items() if len(v) > 1}
+
+
+def get_anchors_and_aliases(filepath):
+    anchors = defaultdict(list)
+    current_level = {}
+    anchor_to_path = {}
+
+    with open(filepath, "r") as file:
+        for line in file:
+            indent = len(line) - len(line.lstrip(" "))
+            key_match = re.search(r"(\w+):", line)
+            anchor_match = re.search(r"&(\w+)", line)
+            alias_match = re.search(r"\*(\w+)", line)
+            if key_match:
+                key = key_match.group(1)
+                # Compute the full path of the current key.
+                full_path = ".".join(
+                    [current_level[i] for i in sorted(current_level.keys()) if i < indent] + [key]
+                )
+                current_level[indent] = key
+                # Remove any keys that are indented more than the current line.
+                keys_to_remove = [i for i in current_level if i > indent]
+                for i in keys_to_remove:
+                    del current_level[i]
+            else:
+                full_path = ".".join([current_level[i] for i in sorted(current_level.keys())])
+            if anchor_match:
+                anchor = anchor_match.group(1)
+                anchor_to_path[anchor] = full_path
+            if alias_match:
+                alias = alias_match.group(1)
+                if alias in anchor_to_path:
+                    anchors[anchor_to_path[alias]].append(full_path)
+    return anchors
+
+
+refs = get_anchors_and_aliases(CONFIG_FILE)
+import ipdb
+
+ipdb.set_trace()
+pass
diff --git a/goli/data/make_data_splits/make_train_val_test_splits_large.ipynb b/goli/data/make_data_splits/make_train_val_test_splits_large.ipynb
new file mode 100644
index 000000000..ec3e870f4
--- /dev/null
+++ b/goli/data/make_data_splits/make_train_val_test_splits_large.ipynb
@@ -0,0 +1,491 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from typing import List, Tuple, Dict\n",
+    "\n",
+    "import random\n",
+    "import pandas as pd\n",
+    "import numpy as np\n",
+    "import torch\n",
+    "from os.path import join\n",
+    "import datamol as dm"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "BASE_PATH = \"../neurips2023/large-dataset/\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "df_L1000_VCAP.shape (15220, 984)\n",
+      "df_L1000_MCF7.shape (11622, 984)\n",
+      "df_PCBA.shape (1563664, 1332)\n",
+      "df_PCQM4M.shape (3810323, 31)\n"
+     ]
+    }
+   ],
+   "source": [
+    "df_L1000_VCAP = pd.read_csv(join(BASE_PATH, \"LINCS_L1000_VCAP_0-4.csv.gz\"))\n",
+    "print(\"df_L1000_VCAP.shape\", df_L1000_VCAP.shape)\n",
+    "df_L1000_MCF7 = pd.read_csv(join(BASE_PATH, \"LINCS_L1000_MCF7_0-4.csv.gz\"))\n",
+    "print(\"df_L1000_MCF7.shape\", df_L1000_MCF7.shape)\n",
+    "df_PCBA = pd.read_parquet(join(BASE_PATH, \"PCBA_1328_1564k.parquet\"))\n",
+    "print(\"df_PCBA.shape\", df_PCBA.shape)\n",
+    "df_PCQM4M = pd.read_parquet(join(BASE_PATH, \"PCQM4M_G25_N4.parquet\"))\n",
+    "print(\"df_PCQM4M.shape\", df_PCQM4M.shape)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def smiles_to_unique_ids(smiles: List):\n",
+    "    return dm.parallelized_with_batches(loop_smiles_to_unique_ids, smiles, batch_size=100, n_jobs=32, progress=True)\n",
+    "\n",
+    "def loop_smiles_to_unique_ids(smiles: List):\n",
+    "    unique_ids = []\n",
+    "    for s in smiles:\n",
+    "        if not isinstance(s, str):\n",
+    "            unique_ids.append(None)\n",
+    "            continue\n",
+    "        mol = dm.to_mol(s)\n",
+    "        if mol is None:\n",
+    "            unique_ids.append(None)\n",
+    "        else:\n",
+    "            unique_ids.append(dm.unique_id(mol))\n",
+    "    return unique_ids"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "a1bdefc329d14d288510253cc36867a8",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/38103 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "[16:03:00] WARNING: not removing hydrogen atom without neighbors\n",
+      "[16:03:10] WARNING: not removing hydrogen atom without neighbors\n",
+      "[16:03:37] WARNING: not removing hydrogen atom without neighbors\n",
+      "[16:03:37] WARNING: not removing hydrogen atom without neighbors\n",
+      "[16:03:37] WARNING: not removing hydrogen atom without neighbors\n",
+      "[16:03:44] WARNING: not removing hydrogen atom without neighbors\n",
+      "[16:03:44] WARNING: not removing hydrogen atom without neighbors\n",
+      "[16:03:44] WARNING: not removing hydrogen atom without neighbors\n",
+      "[16:03:44] WARNING: not removing hydrogen atom without neighbors\n",
+      "[16:03:44] WARNING: not removing hydrogen atom without neighbors\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "7087b7c562ca46ce837695034e38cb57",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/152 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "[16:03:49] SMILES Parse Error: syntax error while parsing: restricted\n",
+      "[16:03:49] SMILES Parse Error: Failed parsing SMILES 'restricted' for input: 'restricted'\n",
+      "[16:03:49] SMILES Parse Error: syntax error while parsing: restricted\n",
+      "[16:03:49] SMILES Parse Error: Failed parsing SMILES 'restricted' for input: 'restricted'\n",
+      "[16:03:49] SMILES Parse Error: syntax error while parsing: restricted\n",
+      "[16:03:49] SMILES Parse Error: Failed parsing SMILES 'restricted' for input: 'restricted'\n",
+      "[16:03:49] SMILES Parse Error: syntax error while parsing: restricted\n",
+      "[16:03:49] SMILES Parse Error: Failed parsing SMILES 'restricted' for input: 'restricted'\n",
+      "[16:03:49] SMILES Parse Error: syntax error while parsing: restricted\n",
+      "[16:03:49] SMILES Parse Error: Failed parsing SMILES 'restricted' for input: 'restricted'\n",
+      "[16:03:49] SMILES Parse Error: syntax error while parsing: restricted\n",
+      "[16:03:49] SMILES Parse Error: Failed parsing SMILES 'restricted' for input: 'restricted'\n",
+      "[16:03:49] SMILES Parse Error: syntax error while parsing: restricted\n",
+      "[16:03:49] SMILES Parse Error: Failed parsing SMILES 'restricted' for input: 'restricted'\n",
+      "[16:03:49] SMILES Parse Error: syntax error while parsing: restricted\n",
+      "[16:03:49] SMILES Parse Error: Failed parsing SMILES 'restricted' for input: 'restricted'\n",
+      "[16:03:49] SMILES Parse Error: syntax error while parsing: restricted\n",
+      "[16:03:49] SMILES Parse Error: Failed parsing SMILES 'restricted'[16:03:49]  for input: 'restricted'\n",
+      "SMILES Parse Error: syntax error while parsing: restricted\n",
+      "[16:03:49] SMILES Parse Error: Failed parsing SMILES 'restricted' for input: 'restricted'\n",
+      "[16:03:49] SMILES Parse Error: syntax error while parsing: restricted\n",
+      "[16:03:49] SMILES Parse Error: Failed parsing SMILES 'restricted' for input: 'restricted'\n",
+      "[16:03:49] SMILES Parse Error: syntax error while parsing: restricted\n",
+      "[16:03:49] SMILES Parse Error: Failed parsing SMILES 'restricted' for input: 'restricted'\n",
+      "[16:03:49] SMILES Parse Error: syntax error while parsing: restricted\n",
+      "[16:03:49] SMILES Parse Error: Failed parsing SMILES 'restricted' for input: 'restricted'\n",
+      "[16:03:49] SMILES Parse Error: syntax error while parsing: restricted\n",
+      "[16:03:49] SMILES Parse Error: Failed parsing SMILES 'restricted' for input: 'restricted'\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "f651dcb510004bb7afef0b02fb1e3b20",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/116 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "[16:03:49] SMILES Parse Error: syntax error while parsing: restricted\n",
+      "[16:03:49] SMILES Parse Error: Failed parsing SMILES 'restricted' for input: 'restricted'\n",
+      "[16:03:49] SMILES Parse Error: syntax error while parsing: restricted\n",
+      "[16:03:49] SMILES Parse Error: Failed parsing SMILES 'restricted' for input: 'restricted'\n",
+      "[16:03:49] SMILES Parse Error: syntax error while parsing: restricted\n",
+      "[16:03:49] SMILES Parse Error: Failed parsing SMILES 'restricted' for input: 'restricted'\n",
+      "[16:03:49] SMILES Parse Error: syntax error while parsing: restricted\n",
+      "[16:03:49] SMILES Parse Error: Failed parsing SMILES 'restricted' for input: 'restricted'\n",
+      "[16:03:49] SMILES Parse Error: syntax error while parsing: restricted\n",
+      "[16:03:49] SMILES Parse Error: Failed parsing SMILES 'restricted' for input: 'restricted'\n",
+      "[16:03:49] SMILES Parse Error: syntax error while parsing: restricted\n",
+      "[16:03:49] SMILES Parse Error: Failed parsing SMILES 'restricted' for input: 'restricted'\n",
+      "[16:03:49] SMILES Parse Error: syntax error while parsing: restricted\n",
+      "[16:03:49] SMILES Parse Error: Failed parsing SMILES 'restricted' for input: 'restricted'\n",
+      "[16:03:49] SMILES Parse Error: syntax error while parsing: restricted\n",
+      "[16:03:49] SMILES Parse Error: Failed parsing SMILES 'restricted' for input: 'restricted'\n",
+      "[16:03:49] SMILES Parse Error: syntax error while parsing: restricted\n",
+      "[16:03:49] SMILES Parse Error: Failed parsing SMILES 'restricted' for input: 'restricted'\n",
+      "[16:03:49] SMILES Parse Error: syntax error while parsing: restricted\n",
+      "[16:03:49] SMILES Parse Error: Failed parsing SMILES 'restricted' for input: 'restricted'\n",
+      "[16:03:49] SMILES Parse Error: syntax error while parsing: restricted\n",
+      "[16:03:49] SMILES Parse Error: Failed parsing SMILES 'restricted' for input: 'restricted'\n",
+      "[16:03:49] SMILES Parse Error: syntax error while parsing: restricted\n",
+      "[16:03:49] SMILES Parse Error: Failed parsing SMILES 'restricted' for input: 'restricted'\n",
+      "[16:03:49] SMILES Parse Error: syntax error while parsing: restricted\n",
+      "[16:03:49] SMILES Parse Error: Failed parsing SMILES 'restricted' for input: 'restricted'\n",
+      "[16:03:49] SMILES Parse Error: syntax error while parsing: restricted\n",
+      "[16:03:49] SMILES Parse Error: Failed parsing SMILES 'restricted' for input: 'restricted'\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "1ab2cca58d1f4a36ac5994b897cfe940",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/15636 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "[16:03:51] WARNING: not removing hydrogen atom without neighbors\n",
+      "[16:03:52] WARNING: not removing hydrogen atom without neighbors\n",
+      "[16:03:53] WARNING: not removing hydrogen atom without neighbors\n",
+      "[16:03:55] WARNING: not removing hydrogen atom without neighbors\n",
+      "[16:03:59] WARNING: not removing hydrogen atom without neighbors\n",
+      "[16:03:59] WARNING: not removing hydrogen atom without neighbors\n",
+      "[16:03:59] WARNING: not removing hydrogen atom without neighbors\n",
+      "[16:04:02] WARNING: not removing hydrogen atom without neighbors\n",
+      "[16:04:02] WARNING: not removing hydrogen atom without neighbors\n",
+      "[16:04:07] WARNING: not removing hydrogen atom without neighbors\n",
+      "[16:04:08] WARNING: not removing hydrogen atom without neighbors\n",
+      "[16:04:08] WARNING: not removing hydrogen atom without neighbors\n",
+      "[16:04:08] WARNING: not removing hydrogen atom without neighbors\n",
+      "[16:04:08] WARNING: not removing hydrogen atom without neighbors\n",
+      "[16:04:09] WARNING: not removing hydrogen atom without neighbors\n",
+      "[16:04:12] WARNING: not removing hydrogen atom without neighbors\n",
+      "[16:04:13] WARNING: not removing hydrogen atom without neighbors\n",
+      "[16:04:13] WARNING: not removing hydrogen atom without neighbors\n",
+      "[16:04:13] WARNING: not removing hydrogen atom without neighbors\n",
+      "[16:04:13] WARNING: not removing hydrogen atom without neighbors\n",
+      "[16:04:13] WARNING: not removing hydrogen atom without neighbors\n",
+      "[16:04:13] WARNING: not removing hydrogen atom without neighbors\n",
+      "[16:04:15] WARNING: not removing hydrogen atom without neighbors\n",
+      "[16:04:15] WARNING: not removing hydrogen atom without neighbors\n",
+      "[16:04:15] WARNING: not removing hydrogen atom without neighbors\n",
+      "[16:04:15] WARNING: not removing hydrogen atom without neighbors\n",
+      "[16:04:15] WARNING: not removing hydrogen atom without neighbors\n",
+      "[16:04:15] WARNING: not removing hydrogen atom without neighbors\n",
+      "[16:04:15] WARNING: not removing hydrogen atom without neighbors\n",
+      "[16:04:15] WARNING: not removing hydrogen atom without neighbors\n",
+      "[16:04:15] WARNING: not removing hydrogen atom without neighbors\n",
+      "[16:04:15] WARNING: not removing hydrogen atom without neighbors\n",
+      "[16:04:16] WARNING: not removing hydrogen atom without neighbors\n",
+      "[16:04:20] WARNING: not removing hydrogen atom without neighbors\n",
+      "[16:04:20] WARNING: not removing hydrogen atom without neighbors\n",
+      "[16:04:22] WARNING: not removing hydrogen atom without neighbors\n",
+      "[16:04:22] WARNING: not removing hydrogen atom without neighbors\n"
+     ]
+    }
+   ],
+   "source": [
+    "unique_ids_QM = smiles_to_unique_ids(df_PCQM4M[\"ordered_smiles\"])\n",
+    "unique_ids_L1000_VCAP = smiles_to_unique_ids(df_L1000_VCAP[\"SMILES\"])\n",
+    "unique_ids_L1000_MCF7 = smiles_to_unique_ids(df_L1000_MCF7[\"SMILES\"])\n",
+    "unique_ids_PCBA = smiles_to_unique_ids(df_PCBA[\"SMILES\"])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "L1000_VCAP 726\n",
+      "L1000_MCF7 1023\n",
+      "PCBA 56512\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Check the number of unique ids that intersect between unique_ids_QM and the other columns\n",
+    "intersection_VCAP = set(unique_ids_QM) & set(unique_ids_L1000_VCAP)\n",
+    "print(\"L1000_VCAP\", len(intersection_VCAP))\n",
+    "\n",
+    "intersection_MCF7 = set(unique_ids_QM) & set(unique_ids_L1000_MCF7)\n",
+    "print(\"L1000_MCF7\", len(intersection_MCF7))\n",
+    "\n",
+    "intersection_PCBA = set(unique_ids_QM) & set(unique_ids_PCBA)\n",
+    "print(\"PCBA\", len(intersection_PCBA))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "QM_in_VCAP 1443 VCAP_in_QM 748\n",
+      "QM_in_MCF7 2373 MCF7_in_QM 1065\n",
+      "QM_in_PCBA 91174 PCBA_in_QM 56512\n"
+     ]
+    }
+   ],
+   "source": [
+    "def find_indices(list_1, list_2):\n",
+    "    intersection = set(list_1) & set(list_2)\n",
+    "    intersection = {elem for elem in intersection if elem is not None}\n",
+    "    is_2_in_1 = np.isin(list_2, list(intersection))\n",
+    "    is_1_in_2 = np.isin(list_1, list(intersection))\n",
+    "    return is_1_in_2, is_2_in_1\n",
+    "\n",
+    "QM_in_VCAP, VCAP_in_QM = find_indices(unique_ids_QM, unique_ids_L1000_VCAP)\n",
+    "print(\"QM_in_VCAP\", sum(QM_in_VCAP), \"VCAP_in_QM\", sum(VCAP_in_QM))\n",
+    "\n",
+    "QM_in_MCF7, MCF7_in_QM = find_indices(unique_ids_QM, unique_ids_L1000_MCF7)\n",
+    "print(\"QM_in_MCF7\", sum(QM_in_MCF7), \"MCF7_in_QM\", sum(MCF7_in_QM))\n",
+    "\n",
+    "QM_in_PCBA, PCBA_in_QM = find_indices(unique_ids_QM, unique_ids_PCBA)\n",
+    "print(\"QM_in_PCBA\", sum(QM_in_PCBA), \"PCBA_in_QM\", sum(PCBA_in_QM))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "test_seen_VCAP = np.where(VCAP_in_QM)[0].tolist()\n",
+    "test_seen_MCF7 = np.where(MCF7_in_QM)[0].tolist()\n",
+    "test_seen_PCBA = np.where(PCBA_in_QM)[0].tolist()\n",
+    "\n",
+    "train_QM_seen = np.where(QM_in_VCAP | QM_in_MCF7 | QM_in_PCBA)[0].tolist()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def make_random_splits(num_elem, val_ratio, test_ratio, seed=42, ignore_idx=None):\n",
+    "    \"\"\"Make random splits of the data into train, validation, and test sets.\n",
+    "\n",
+    "    Args:\n",
+    "        num_elem (int): Number of elements in the dataset.\n",
+    "        val_ratio (float): Ratio of the dataset to use for validation.\n",
+    "        test_ratio (float): Ratio of the dataset to use for testing.\n",
+    "        seed (int): Random seed.\n",
+    "        ignore_idx (list): List of indices to ignore.\n",
+    "\n",
+    "    Returns:\n",
+    "        train_idx (list): List of indices for the training set.\n",
+    "        val_idx (list): List of indices for the validation set.\n",
+    "        test_idx (list): List of indices for the test set.\n",
+    "    \"\"\"\n",
+    "    # Create a list of indices\n",
+    "    idx = list(range(num_elem))\n",
+    "    # Remove the indices to ignore\n",
+    "    if ignore_idx is not None:\n",
+    "        idx = list(set(idx) - set(ignore_idx))\n",
+    "        num_elem = len(idx)\n",
+    "    # Shuffle the list of indices\n",
+    "    random.seed(seed)\n",
+    "    random.shuffle(idx)\n",
+    "    # Compute the number of elements in each set\n",
+    "    num_val = int(num_elem * val_ratio)\n",
+    "    num_test = int(num_elem * test_ratio)\n",
+    "    num_train = num_elem - num_val - num_test\n",
+    "    # Split the list of indices into three sets\n",
+    "    train_idx = idx[:num_train]\n",
+    "    val_idx = idx[num_train:(num_train + num_val)]\n",
+    "    test_idx = idx[(num_train + num_val):]\n",
+    "    # Return the three lists of indices\n",
+    "    return train_idx, val_idx, test_idx"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "train_VCAP, val_VCAP, test_VCAP = make_random_splits(len(df_L1000_VCAP), 0.04, 0.04, seed=42, ignore_idx=test_seen_VCAP)\n",
+    "train_MCF7, val_MCF7, test_MCF7 = make_random_splits(len(df_L1000_MCF7), 0.04, 0.04, seed=42, ignore_idx=test_seen_MCF7)\n",
+    "train_PCBA, val_PCBA, test_PCBA = make_random_splits(len(df_PCBA), 0.04, 0.04, seed=42, ignore_idx=test_seen_PCBA)\n",
+    "train_QM, val_QM, test_QM = make_random_splits(len(df_PCQM4M), 0.04, 0.04, seed=42, ignore_idx=train_QM_seen)\n",
+    "train_QM = np.concatenate((train_QM, train_QM_seen))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def make_random_splits_file(df, out_file, train_idx, val_idx, test_idx, test_seen_idx=None):\n",
+    "    # Save the splits\n",
+    "    if test_seen_idx is None:\n",
+    "        splits_dict = {\"train\": train_idx, \"val\": val_idx, \"test\": test_idx}\n",
+    "    else:\n",
+    "        splits_dict = {\"train\": train_idx, \"val\": val_idx, \"test\": test_idx, \"test_seen\": test_seen_idx}\n",
+    "\n",
+    "    # Check the splits validity\n",
+    "    \n",
+    "    assert len(set(train_idx).intersection(set(val_idx))) == 0\n",
+    "    assert len(set(train_idx).intersection(set(test_idx))) == 0\n",
+    "    assert len(set(val_idx).intersection(set(test_idx))) == 0\n",
+    "    assert len(train_idx) > 0\n",
+    "    assert len(val_idx) > 0\n",
+    "    assert len(test_idx) > 0\n",
+    "\n",
+    "    if test_seen_idx is None:\n",
+    "        assert len(df) == len(train_idx) + len(val_idx) + len(test_idx), f\"{len(df)} != {len(train_idx)} + {len(val_idx)} + {len(test_idx)}\"\n",
+    "        print(out_file, \"train\", len(train_idx), \"val\", len(val_idx), \"test\", len(test_idx))\n",
+    "    else:\n",
+    "        assert len(test_seen_idx) > 0\n",
+    "        assert len(set(train_idx).intersection(set(test_seen_idx))) == 0\n",
+    "        assert len(set(val_idx).intersection(set(test_seen_idx))) == 0\n",
+    "        assert len(set(test_idx).intersection(set(test_seen_idx))) == 0\n",
+    "        assert len(df) == len(train_idx) + len(val_idx) + len(test_idx) + len(test_seen_idx), f\"{len(df)} != {len(train_idx)} + {len(val_idx)} + {len(test_idx)} + {len(test_seen_idx)}\"\n",
+    "        print(out_file, \"train\", len(train_idx), \"val\", len(val_idx), \"test\", len(test_idx), \"test_seen\", len(test_seen_idx))\n",
+    "\n",
+    "    torch.save(splits_dict, out_file)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "l1000_vcap_random_splits.pt train 13316 val 578 test 578 test_seen 748\n",
+      "l1000_mcf7_random_splits.pt train 9713 val 422 test 422 test_seen 1065\n",
+      "pcba_1328_random_splits.pt train 1386580 val 60286 test 60286 test_seen 56512\n",
+      "pcqm4m_g25_n4_random_splits.pt train 3512805 val 148759 test 148759\n"
+     ]
+    }
+   ],
+   "source": [
+    "make_random_splits_file(df_L1000_VCAP, \"l1000_vcap_random_splits.pt\", train_VCAP, val_VCAP, test_VCAP, test_seen_VCAP)\n",
+    "make_random_splits_file(df_L1000_MCF7, \"l1000_mcf7_random_splits.pt\", train_MCF7, val_MCF7, test_MCF7, test_seen_MCF7)\n",
+    "make_random_splits_file(df_PCBA, \"pcba_1328_random_splits.pt\", train_PCBA, val_PCBA, test_PCBA, test_seen_PCBA)\n",
+    "make_random_splits_file(df_PCQM4M, \"pcqm4m_g25_n4_random_splits.pt\", train_QM, val_QM, test_QM)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "goli",
+   "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.8.10"
+  },
+  "orig_nbformat": 4
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}
diff --git a/goli/data/make_data_splits/make_train_val_test_splits_small.ipynb b/goli/data/make_data_splits/make_train_val_test_splits_small.ipynb
new file mode 100644
index 000000000..f0c523325
--- /dev/null
+++ b/goli/data/make_data_splits/make_train_val_test_splits_small.ipynb
@@ -0,0 +1,106 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import random\n",
+    "import pandas as pd\n",
+    "import torch"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def make_random_splits(num_elem, val_ratio, test_ratio, seed=42):\n",
+    "    \"\"\"Make random splits of the data into train, validation, and test sets.\n",
+    "\n",
+    "    Args:\n",
+    "        num_elem (int): Number of elements in the dataset.\n",
+    "        val_ratio (float): Ratio of the dataset to use for validation.\n",
+    "        test_ratio (float): Ratio of the dataset to use for testing.\n",
+    "\n",
+    "    Returns:\n",
+    "        train_idx (list): List of indices for the training set.\n",
+    "        val_idx (list): List of indices for the validation set.\n",
+    "        test_idx (list): List of indices for the test set.\n",
+    "    \"\"\"\n",
+    "    # Create a list of indices\n",
+    "    idx = list(range(num_elem))\n",
+    "    # Shuffle the list of indices\n",
+    "    random.seed(seed)\n",
+    "    random.shuffle(idx)\n",
+    "    # Compute the number of elements in each set\n",
+    "    num_val = int(num_elem * val_ratio)\n",
+    "    num_test = int(num_elem * test_ratio)\n",
+    "    num_train = num_elem - num_val - num_test\n",
+    "    # Split the list of indices into three sets\n",
+    "    train_idx = idx[:num_train]\n",
+    "    val_idx = idx[num_train:num_train + num_val]\n",
+    "    test_idx = idx[num_train + num_val:]\n",
+    "    # Return the three lists of indices\n",
+    "    return train_idx, val_idx, test_idx"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def make_random_splits_from_file(in_file, out_file, val_ratio, test_ratio, seed=42):\n",
+    "    # Make the splits for QM9\n",
+    "    df = pd.read_csv(in_file, usecols=[\"smiles\"])\n",
+    "    train_idx, val_idx, test_idx = make_random_splits(len(df), val_ratio, test_ratio, seed=seed)\n",
+    "    # Save the splits\n",
+    "    splits_dict = {\"train\": train_idx, \"val\": val_idx, \"test\": test_idx}\n",
+    "\n",
+    "    # Check the splits validity\n",
+    "    assert len(set(train_idx).intersection(set(val_idx))) == 0\n",
+    "    assert len(set(train_idx).intersection(set(test_idx))) == 0\n",
+    "    assert len(set(val_idx).intersection(set(test_idx))) == 0\n",
+    "    assert len(train_idx) + len(val_idx) + len(test_idx) == len(df)\n",
+    "    \n",
+    "    torch.save(splits_dict, out_file)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "make_random_splits_from_file(\"qm9.csv.gz\", \"qm9_random_splits.pt\", 0.1, 0.1)\n",
+    "make_random_splits_from_file(\"Tox21-7k-12-labels.csv.gz\", \"Tox21_random_splits.pt\", 0.1, 0.1)\n",
+    "make_random_splits_from_file(\"ZINC12k.csv.gz\", \"ZINC12k_random_splits.pt\", 0.1, 0.1)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "goli",
+   "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.9"
+  },
+  "orig_nbformat": 4
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}
diff --git a/tests/.gitignore b/tests/.gitignore
new file mode 100644
index 000000000..9262b8229
--- /dev/null
+++ b/tests/.gitignore
@@ -0,0 +1,1 @@
+temp_cache_*
\ No newline at end of file
diff --git a/tests/converted_fake_multilevel_data.parquet b/tests/converted_fake_multilevel_data.parquet
new file mode 100644
index 000000000..53c437fe8
Binary files /dev/null and b/tests/converted_fake_multilevel_data.parquet differ
diff --git a/tests/fake_and_missing_multilevel_data.parquet b/tests/fake_and_missing_multilevel_data.parquet
new file mode 100644
index 000000000..27624cf8b
Binary files /dev/null and b/tests/fake_and_missing_multilevel_data.parquet differ
diff --git a/tests/test_base_layers.py b/tests/test_base_layers.py
index 287685e14..e9d861ba0 100644
--- a/tests/test_base_layers.py
+++ b/tests/test_base_layers.py
@@ -8,6 +8,7 @@
 from copy import deepcopy
 
 from goli.nn.base_layers import DropPath, TransformerEncoderLayerMup
+from goli.ipu.to_dense_batch import to_dense_batch, to_sparse_batch
 
 
 class test_Base_Layers(ut.TestCase):
@@ -27,6 +28,9 @@ class test_Base_Layers(ut.TestCase):
     g2 = Data(feat=x2, edge_index=edge_idx2, edge_feat=e2)
     bg = Batch.from_data_list([g1, g2])
 
+    batch_size = 2
+    max_num_nodes_per_graph = max(x1.shape[0], x2.shape[0])
+
     # for drop_rate=0.5, test if the output shape is correct
     def test_droppath_layer_0p5(self):
         bg = deepcopy(self.bg)
@@ -59,5 +63,19 @@ def test_transformer_encoder_layer_mup(self):
         layer = TransformerEncoderLayerMup(
             biased_attention=False, d_model=self.in_dim, nhead=1, dim_feedforward=4 * self.in_dim
         )
-        h_out = layer.forward(feat_in)
+
+        feat_dense, key_padding_mask, idx = 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,
+        )
+        attn_mask = None
+        key_padding_mask = ~key_padding_mask
+
+        h_out_dense = layer.forward(feat_dense)
+
+        h_out = to_sparse_batch(h_out_dense, mask_idx=idx)
+
         self.assertEqual(h_out.shape, feat_in.shape)
diff --git a/tests/test_collate.py b/tests/test_collate.py
index cc96dfd1c..fabba5aa9 100644
--- a/tests/test_collate.py
+++ b/tests/test_collate.py
@@ -6,6 +6,7 @@
 from copy import deepcopy
 import torch
 import numpy as np
+from torch_geometric.data import Data
 
 from goli.data.collate import collate_labels, goli_collate_fn
 
@@ -13,62 +14,123 @@
 class test_Collate(ut.TestCase):
     def test_collate_labels(self):
         # Create fake labels
-        labels_size_dict = {"label1": [1], "label2": [5], "label3": [5, 2]}
+        labels_size_dict = {
+            "graph_label1": [1],
+            "graph_label2": [3],
+            "node_label2": [5],
+            "edge_label3": [5, 2],
+            "node_label4": [5, 1],
+        }
         fake_label = {
-            "label1": torch.FloatTensor([1]),
-            "label2": torch.FloatTensor([1, 2, 3, 4, 5]),
-            "label3": torch.FloatTensor([[1, 2], [3, 4], [5, 6], [7, 8], [9, 10]]),
+            "graph_label1": torch.FloatTensor([1]),
+            "graph_label2": torch.FloatTensor([1, 2, 3]),
+            "node_label2": torch.FloatTensor([1, 2, 3, 4, 5]),
+            "edge_label3": torch.FloatTensor([[1, 2], [3, 4], [5, 6], [7, 8], [9, 10]]),
+            "node_label4": torch.FloatTensor([[1], [2], [3], [4], [5]]),
         }
         fake_labels = []
         num_labels = 10
         for i in range(num_labels):
-            fake_labels.append(deepcopy({key: val + 17 * i for key, val in fake_label.items()}))
+            pyg_labels = Data(x=torch.empty(5, 0), edge_index=torch.empty(2, 5))
+            for key, val in fake_label.items():
+                pyg_labels[key] = val + 17 * 2
+            fake_labels.append(pyg_labels)
 
         # Collate labels and check for the right shapes
-        collated_labels = collate_labels(deepcopy(fake_labels), labels_size_dict)
-        self.assertEqual(collated_labels["label1"].shape, torch.Size([num_labels, 1]))
-        self.assertEqual(collated_labels["label2"].shape, torch.Size([num_labels, 5]))
-        self.assertEqual(collated_labels["label3"].shape, torch.Size([num_labels, 5, 2]))
+        collated_labels = collate_labels(deepcopy(fake_labels), deepcopy(labels_size_dict))
+        self.assertEqual(collated_labels["graph_label1"].shape, torch.Size([num_labels, 1]))  # , 1
+        self.assertEqual(collated_labels["graph_label2"].shape, torch.Size([num_labels, 3]))  # , 1
+        self.assertEqual(collated_labels["node_label2"].shape, torch.Size([num_labels * 5, 1]))  # , 5
+        self.assertEqual(collated_labels["edge_label3"].shape, torch.Size([num_labels * 5, 2]))  # , 5, 2
+        self.assertEqual(collated_labels["node_label4"].shape, torch.Size([num_labels * 5, 1]))  # , 5, 1
 
         # Check that the values are correct
-        label1_true = deepcopy(torch.stack([this_label["label1"] for this_label in fake_labels]))
-        label2_true = deepcopy(torch.stack([this_label["label2"] for this_label in fake_labels]))
-        label3_true = deepcopy(torch.stack([this_label["label3"] for this_label in fake_labels]))
-        np.testing.assert_array_equal(collated_labels["label1"].numpy(), label1_true.numpy())
-        np.testing.assert_array_equal(collated_labels["label2"].numpy(), label2_true.numpy())
-        np.testing.assert_array_equal(collated_labels["label3"].numpy(), label3_true.numpy())
+        graph_label1_true = deepcopy(torch.stack([this_label["graph_label1"] for this_label in fake_labels]))
+        graph_label2_true = deepcopy(torch.stack([this_label["graph_label2"] for this_label in fake_labels]))
+        label2_true = deepcopy(torch.stack([this_label["node_label2"] for this_label in fake_labels]))
+        label3_true = deepcopy(torch.stack([this_label["edge_label3"] for this_label in fake_labels]))
+        label4_true = deepcopy(torch.stack([this_label["node_label4"] for this_label in fake_labels]))
+
+        # NOTE: Flatten due to the way Data objects are collated (concat along first dim, instead of stacked)
+        np.testing.assert_array_equal(collated_labels["graph_label1"].numpy(), graph_label1_true.numpy())
+        np.testing.assert_array_equal(collated_labels["graph_label2"].numpy(), graph_label2_true.numpy())
+        np.testing.assert_array_equal(
+            collated_labels["node_label2"].numpy(), label2_true.flatten(0, 1).unsqueeze(1).numpy()
+        )
+        np.testing.assert_array_equal(
+            collated_labels["edge_label3"].numpy(), label3_true.flatten(0, 1).numpy()
+        )
+        np.testing.assert_array_equal(
+            collated_labels["node_label4"].numpy(), label4_true.flatten(0, 1).numpy()
+        )
 
         # Remove some labels and check that the collation still works and puts `nan` in the right places
-        missing_labels = {"label1": [1, 3, 5, 7, 9], "label2": [0, 2, 4, 6, 8], "label3": [0, 1, 2, 3, 4]}
+        missing_labels = {
+            "graph_label1": [1, 3, 5, 7, 9],
+            "graph_label2": [0, 4, 3, 1, 7],
+            "node_label2": [0, 2, 4, 6, 8],
+            "edge_label3": [0, 1, 2, 3, 4],
+            "node_label4": [5, 1, 4, 9, 6],
+        }
         for key, missing_idx in missing_labels.items():
             for idx in missing_idx:
                 fake_labels[idx].pop(key)
-        label1_true[missing_labels["label1"]] = float("nan")
-        label2_true[missing_labels["label2"]] = float("nan")
-        label3_true[missing_labels["label3"]] = float("nan")
+        graph_label1_true[missing_labels["graph_label1"]] = float("nan")
+        graph_label2_true[missing_labels["graph_label2"]] = float("nan")
+        label2_true[missing_labels["node_label2"]] = float("nan")
+        label3_true[missing_labels["edge_label3"]] = float("nan")
+        label4_true[missing_labels["node_label4"]] = float("nan")
 
         # Collate labels and check for the right shapes
-        collated_labels = collate_labels(deepcopy(fake_labels), labels_size_dict)
-        self.assertEqual(collated_labels["label1"].shape, torch.Size([num_labels, 1]))
-        self.assertEqual(collated_labels["label2"].shape, torch.Size([num_labels, 5]))
-        self.assertEqual(collated_labels["label3"].shape, torch.Size([num_labels, 5, 2]))
+        labels_size_dict = {
+            "graph_label1": [1],
+            "graph_label2": [3],
+            "node_label2": [5],
+            "edge_label3": [5, 2],
+            "node_label4": [5, 1],
+        }
+        collated_labels = collate_labels(deepcopy(fake_labels), deepcopy(labels_size_dict))
+        self.assertEqual(collated_labels["graph_label1"].shape, torch.Size([num_labels, 1]))  # , 1
+        self.assertEqual(collated_labels["graph_label2"].shape, torch.Size([num_labels, 3]))  # , 1
+        self.assertEqual(collated_labels["node_label2"].shape, torch.Size([num_labels * 5, 1]))  # , 5
+        self.assertEqual(collated_labels["edge_label3"].shape, torch.Size([num_labels * 5, 2]))  # , 5, 2
+        self.assertEqual(collated_labels["node_label4"].shape, torch.Size([num_labels * 5, 1]))  # , 5, 1
 
         # Check that the values are correct when some labels are missing
-        np.testing.assert_array_equal(collated_labels["label1"].numpy(), label1_true.numpy())
-        np.testing.assert_array_equal(collated_labels["label2"].numpy(), label2_true.numpy())
-        np.testing.assert_array_equal(collated_labels["label3"].numpy(), label3_true.numpy())
-
+        # NOTE: Flatten due to the way Data objects are collated (concat along first dim, instead of stacked)
+        np.testing.assert_array_equal(collated_labels["graph_label1"].numpy(), graph_label1_true.numpy())
+        np.testing.assert_array_equal(collated_labels["graph_label2"].numpy(), graph_label2_true.numpy())
+        np.testing.assert_array_equal(
+            collated_labels["node_label2"].numpy(), label2_true.flatten(0, 1).unsqueeze(1).numpy()
+        )
+        np.testing.assert_array_equal(
+            collated_labels["edge_label3"].numpy(), label3_true.flatten(0, 1).numpy()
+        )
+        np.testing.assert_array_equal(
+            collated_labels["node_label4"].numpy(), label4_true.flatten(0, 1).numpy()
+        )
         # Now test the `goli_collate_fn` function when only labels are given
         fake_labels2 = [{"labels": this_label} for this_label in fake_labels]
         collated_labels = goli_collate_fn(deepcopy(fake_labels2), labels_size_dict=labels_size_dict)["labels"]
-        self.assertEqual(collated_labels["label1"].shape, torch.Size([num_labels, 1]))
-        self.assertEqual(collated_labels["label2"].shape, torch.Size([num_labels, 5]))
-        self.assertEqual(collated_labels["label3"].shape, torch.Size([num_labels, 5, 2]))
+        self.assertEqual(collated_labels["graph_label1"].shape, torch.Size([num_labels, 1]))
+        self.assertEqual(collated_labels["graph_label2"].shape, torch.Size([num_labels, 3]))
+        self.assertEqual(collated_labels["node_label2"].shape, torch.Size([num_labels * 5, 1]))  # , 5
+        self.assertEqual(collated_labels["edge_label3"].shape, torch.Size([num_labels * 5, 2]))  # , 5, 2
+        self.assertEqual(collated_labels["node_label4"].shape, torch.Size([num_labels * 5, 1]))  # , 5, 1
 
         # Check that the values are correct when some labels are missing
-        np.testing.assert_array_equal(collated_labels["label1"].numpy(), label1_true.numpy())
-        np.testing.assert_array_equal(collated_labels["label2"].numpy(), label2_true.numpy())
-        np.testing.assert_array_equal(collated_labels["label3"].numpy(), label3_true.numpy())
+        # NOTE: Flatten due to the way Data objects are collated (concat along first dim, instead of stacked)
+        np.testing.assert_array_equal(collated_labels["graph_label1"].numpy(), graph_label1_true.numpy())
+        np.testing.assert_array_equal(collated_labels["graph_label2"].numpy(), graph_label2_true.numpy())
+        np.testing.assert_array_equal(
+            collated_labels["node_label2"].numpy(), label2_true.flatten(0, 1).unsqueeze(1).numpy()
+        )
+        np.testing.assert_array_equal(
+            collated_labels["edge_label3"].numpy(), label3_true.flatten(0, 1).numpy()
+        )
+        np.testing.assert_array_equal(
+            collated_labels["node_label4"].numpy(), label4_true.flatten(0, 1).numpy()
+        )
 
 
 if __name__ == "__main__":
diff --git a/tests/test_datamodule.py b/tests/test_datamodule.py
index eab819cc7..00934285f 100644
--- a/tests/test_datamodule.py
+++ b/tests/test_datamodule.py
@@ -29,7 +29,7 @@ def test_ogb_datamodule(self):
 
         # Config for datamodule
         task_specific_args = {}
-        task_specific_args["task_1"] = {"dataset_name": dataset_name}
+        task_specific_args["task_1"] = {"task_level": "graph", "dataset_name": dataset_name}
         dm_args = {}
         dm_args["cache_data_path"] = None
         dm_args["featurization"] = featurization_args
@@ -61,7 +61,7 @@ def test_ogb_datamodule(self):
         # test batch loader
         batch = next(iter(ds.train_dataloader()))
         assert len(batch["smiles"]) == 16
-        assert len(batch["labels"]["task_1"]) == 16
+        assert len(batch["labels"]["graph_task_1"]) == 16
         assert len(batch["mol_ids"]) == 16
 
     def test_none_filtering(self):
@@ -171,7 +171,7 @@ def test_caching(self):
 
         # Config for datamodule
         task_specific_args = {}
-        task_specific_args["task_1"] = {"dataset_name": dataset_name}
+        task_specific_args["task_1"] = {"task_level": "graph", "dataset_name": dataset_name}
         dm_args = {}
         dm_args["featurization"] = featurization_args
         dm_args["batch_size_training"] = 16
@@ -216,7 +216,7 @@ def test_caching(self):
         # test batch loader
         batch = next(iter(ds.train_dataloader()))
         assert len(batch["smiles"]) == 16
-        assert len(batch["labels"]["task_1"]) == 16
+        assert len(batch["labels"]["graph_task_1"]) == 16
         assert len(batch["mol_ids"]) == 16
 
     def test_datamodule_with_none_molecules(self):
@@ -230,9 +230,24 @@ def test_datamodule_with_none_molecules(self):
         bad_csv = "tests/data/micro_ZINC_corrupt.csv"
         task_specific_args = {}
         task_kwargs = {"df_path": bad_csv, "split_val": 0.0, "split_test": 0.0}
-        task_specific_args["task_1"] = {"label_cols": "SA", "smiles_col": "SMILES1", **task_kwargs}
-        task_specific_args["task_2"] = {"label_cols": "logp", "smiles_col": "SMILES2", **task_kwargs}
-        task_specific_args["task_3"] = {"label_cols": "score", "smiles_col": "SMILES3", **task_kwargs}
+        task_specific_args["task_1"] = {
+            "task_level": "graph",
+            "label_cols": "SA",
+            "smiles_col": "SMILES1",
+            **task_kwargs,
+        }
+        task_specific_args["task_2"] = {
+            "task_level": "graph",
+            "label_cols": "logp",
+            "smiles_col": "SMILES2",
+            **task_kwargs,
+        }
+        task_specific_args["task_3"] = {
+            "task_level": "graph",
+            "label_cols": "score",
+            "smiles_col": "SMILES3",
+            **task_kwargs,
+        }
 
         # Read the corrupted dataset and get stats
         df = pd.read_csv(bad_csv)
@@ -294,7 +309,7 @@ def test_datamodule_with_none_molecules(self):
         train_labels = [{task: val[0] for task, val in d["labels"].items()} for d in datamodule.train_ds]
         train_labels_df = pd.DataFrame(train_labels)
         train_labels_df = train_labels_df.rename(
-            columns={"task_1": "SA", "task_2": "logp", "task_3": "score"}
+            columns={"graph_task_1": "graph_SA", "graph_task_2": "graph_logp", "graph_task_3": "graph_score"}
         )
         train_labels_df["smiles"] = [s[0] for s in datamodule.train_ds.smiles]
         train_labels_df = train_labels_df.set_index("smiles")
@@ -302,7 +317,7 @@ def test_datamodule_with_none_molecules(self):
 
         # Check that the labels are correct
         df2 = df.reset_index()[~bad_smiles].set_index("idx_smiles").sort_index()
-        labels = train_labels_df[["SA", "logp", "score"]].values
+        labels = train_labels_df[["graph_SA", "graph_logp", "graph_score"]].values
         nans = np.isnan(labels)
         true_nans = df2[["SMILES1", "SMILES2", "SMILES3"]].values == "XXX"
         true_labels = df2[["SA", "logp", "score"]].values
@@ -316,7 +331,9 @@ def test_datamodule_multiple_data_files(self):
         # Test single CSV files
         csv_file = "tests/data/micro_ZINC_shard_1.csv"
         task_kwargs = {"df_path": csv_file, "split_val": 0.0, "split_test": 0.0}
-        task_specific_args = {"task": {"label_cols": ["score"], "smiles_col": "SMILES", **task_kwargs}}
+        task_specific_args = {
+            "task": {"task_level": "graph", "label_cols": ["score"], "smiles_col": "SMILES", **task_kwargs}
+        }
 
         ds = MultitaskFromSmilesDataModule(task_specific_args)
         ds.prepare_data()
@@ -327,7 +344,9 @@ def test_datamodule_multiple_data_files(self):
         # Test multi CSV files
         csv_file = "tests/data/micro_ZINC_shard_*.csv"
         task_kwargs = {"df_path": csv_file, "split_val": 0.0, "split_test": 0.0}
-        task_specific_args = {"task": {"label_cols": ["score"], "smiles_col": "SMILES", **task_kwargs}}
+        task_specific_args = {
+            "task": {"task_level": "graph", "label_cols": ["score"], "smiles_col": "SMILES", **task_kwargs}
+        }
 
         ds = MultitaskFromSmilesDataModule(task_specific_args)
         ds.prepare_data()
@@ -338,7 +357,9 @@ def test_datamodule_multiple_data_files(self):
         # Test single Parquet files
         parquet_file = "tests/data/micro_ZINC_shard_1.parquet"
         task_kwargs = {"df_path": parquet_file, "split_val": 0.0, "split_test": 0.0}
-        task_specific_args = {"task": {"label_cols": ["score"], "smiles_col": "SMILES", **task_kwargs}}
+        task_specific_args = {
+            "task": {"task_level": "graph", "label_cols": ["score"], "smiles_col": "SMILES", **task_kwargs}
+        }
 
         ds = MultitaskFromSmilesDataModule(task_specific_args)
         ds.prepare_data()
@@ -349,7 +370,9 @@ def test_datamodule_multiple_data_files(self):
         # Test multi Parquet files
         parquet_file = "tests/data/micro_ZINC_shard_*.parquet"
         task_kwargs = {"df_path": parquet_file, "split_val": 0.0, "split_test": 0.0}
-        task_specific_args = {"task": {"label_cols": ["score"], "smiles_col": "SMILES", **task_kwargs}}
+        task_specific_args = {
+            "task": {"task_level": "graph", "label_cols": ["score"], "smiles_col": "SMILES", **task_kwargs}
+        }
 
         ds = MultitaskFromSmilesDataModule(task_specific_args)
         ds.prepare_data()
diff --git a/tests/test_dataset.py b/tests/test_dataset.py
index 0498e83a2..35b53444c 100644
--- a/tests/test_dataset.py
+++ b/tests/test_dataset.py
@@ -140,16 +140,16 @@ def test_multitask_dataset_case_2(self):
 
             if task == "SA":
                 self.assertEqual(label_SA, multitask_microzinc.labels[found_idx]["SA"])
-                self.assertFalse("score" in multitask_microzinc.labels[found_idx].keys())
-                self.assertFalse("logp" in multitask_microzinc.labels[found_idx].keys())
+                self.assertFalse("score" in multitask_microzinc.labels[found_idx].keys)
+                self.assertFalse("logp" in multitask_microzinc.labels[found_idx].keys)
             elif task == "logp":
                 self.assertEqual(label_logp, multitask_microzinc.labels[found_idx]["logp"])
-                self.assertFalse("score" in multitask_microzinc.labels[found_idx].keys())
-                self.assertFalse("SA" in multitask_microzinc.labels[found_idx].keys())
+                self.assertFalse("score" in multitask_microzinc.labels[found_idx].keys)
+                self.assertFalse("SA" in multitask_microzinc.labels[found_idx].keys)
             elif task == "score":
                 self.assertEqual(label_score, multitask_microzinc.labels[found_idx]["score"])
-                self.assertFalse("SA" in multitask_microzinc.labels[found_idx].keys())
-                self.assertFalse("logp" in multitask_microzinc.labels[found_idx].keys())
+                self.assertFalse("SA" in multitask_microzinc.labels[found_idx].keys)
+                self.assertFalse("logp" in multitask_microzinc.labels[found_idx].keys)
 
     def test_multitask_dataset_case_3(self):
         """Case: Different tasks, but with semi-intersection (some smiles unique per task, some intersect)
diff --git a/tests/test_ipu_losses.py b/tests/test_ipu_losses.py
index 05bcd3e8e..48f93ec4b 100644
--- a/tests/test_ipu_losses.py
+++ b/tests/test_ipu_losses.py
@@ -93,17 +93,17 @@ def test_l1(self):
         # Regular loss
         loss_true = L1Loss()(preds, target)
         loss_ipu = L1LossIPU()(preds, target)
-        self.assertFalse(loss_true.isnan(), "Regular MSELoss is NaN")
+        self.assertFalse(loss_true.isnan(), "Regular MAELoss is NaN")
         self.assertAlmostEqual(
-            loss_true.item(), loss_ipu.item(), places=6, msg="Regular MSELoss is different"
+            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 MSELoss with target_nan is NaN")
-        self.assertFalse(loss_ipu.isnan(), "Regular MSELossIPU with target_nan is 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 MSELoss with NaN is different"
+            loss_true.item(), loss_ipu.item(), places=6, msg="Regular MAELoss with NaN is different"
         )
diff --git a/tests/test_ipu_metrics.py b/tests/test_ipu_metrics.py
index dbaaab8ef..f450f917a 100644
--- a/tests/test_ipu_metrics.py
+++ b/tests/test_ipu_metrics.py
@@ -59,8 +59,8 @@ def test_auroc(self):
         target_nan[target_nan >= 0.5] = 1
 
         # Regular loss
-        score_true = auroc(preds, target.to(int), num_classes=1)
-        score_ipu = auroc_ipu(preds, target, num_classes=1)
+        score_true = auroc(preds, target.to(int))
+        score_ipu = auroc_ipu(preds, target)
         self.assertFalse(score_true.isnan(), "Regular AUROC score is NaN")
         self.assertAlmostEqual(
             score_true.item(), score_ipu.item(), places=6, msg="Regular AUROC score is different"
@@ -68,15 +68,15 @@ def test_auroc(self):
 
         # Weighted loss (As in BCE)
         sample_weights = torch.rand(preds.shape[0], dtype=torch.float32)
-        score_true = auroc(preds, target.to(int), num_classes=1, sample_weights=sample_weights)
-        score_ipu = auroc_ipu(preds, target, num_classes=1, sample_weights=sample_weights)
+        score_true = auroc(preds, target.to(int), sample_weights=sample_weights)
+        score_ipu = auroc_ipu(preds, target, sample_weights=sample_weights)
         self.assertFalse(score_true.isnan(), "Regular AUROC score is NaN")
         self.assertAlmostEqual(score_true.item(), score_ipu.item(), msg="Weighted AUROC score is different")
 
         # Regular loss with NaNs in target
         not_nan = ~target_nan.isnan()
-        score_true = auroc(preds[not_nan], target[not_nan].to(int), num_classes=1)
-        score_ipu = auroc_ipu(preds, target_nan, num_classes=1)
+        score_true = auroc(preds[not_nan], target[not_nan].to(int))
+        score_ipu = auroc_ipu(preds, target_nan)
         self.assertFalse(score_true.isnan(), "Regular AUROC score with target_nan is NaN")
         self.assertFalse(score_ipu.isnan(), "Regular AUROCIPU score with target_nan is NaN")
         self.assertAlmostEqual(
@@ -110,8 +110,8 @@ def test_average_precision(self):  # TODO: Make work with multi-class
         target_nan[target_nan >= 0.5] = 1
 
         # Regular loss
-        score_true = average_precision(preds, target.to(int), num_classes=1)
-        score_ipu = average_precision_ipu(preds, target, num_classes=1)
+        score_true = average_precision(preds, target.to(int), task="binary")
+        score_ipu = average_precision_ipu(preds, target.to(int), task="binary")
         self.assertFalse(score_true.isnan(), "Regular Average Precision is NaN")
         self.assertAlmostEqual(
             score_true.item(), score_ipu.item(), places=6, msg="Regular Average Precision is different"
@@ -119,8 +119,8 @@ def test_average_precision(self):  # TODO: Make work with multi-class
 
         # Regular average precision with NaNs in target
         not_nan = ~target_nan.isnan()
-        score_true = average_precision(preds[not_nan], target[not_nan].to(int), num_classes=1)
-        score_ipu = average_precision_ipu(preds, target_nan, num_classes=1)
+        score_true = average_precision(preds[not_nan], target[not_nan].to(int), task="binary")
+        score_ipu = average_precision_ipu(preds, target_nan, task="binary")
         self.assertFalse(score_true.isnan(), "Regular Average Precision with target_nan is NaN")
         self.assertFalse(score_ipu.isnan(), "Regular Average Precision IPU score with target_nan is NaN")
         self.assertAlmostEqual(
diff --git a/tests/test_metrics.py b/tests/test_metrics.py
index 9b251540f..48f9b447b 100644
--- a/tests/test_metrics.py
+++ b/tests/test_metrics.py
@@ -191,10 +191,10 @@ def test_pickling(self):
                                             self.assertIsInstance(state["metric"], str, msg=err_msg)
 
     def test_classifigression_target_squeezing(self):
-        preds = torch.Tensor([[0.1, 0.1, 0.3, 0.5, 0.0], [0.1, 0.0, 0.7, 0.2, 0.0]])
+        preds = torch.Tensor([[0.1, 0.1, 0.3, 0.5, 0.0, 0.1, 0.0, 0.7, 0.2, 0.0]])
         target = torch.Tensor([3, 0])
         expected_scores = [0.5, 0.75]
-        n_brackets = preds.shape[1]
+        n_brackets = 5
         metrics = ["accuracy", "averageprecision"]
         other_kwargs = [
             {"task": "multiclass", "num_classes": n_brackets, "top_k": 1},
@@ -204,6 +204,7 @@ def test_classifigression_target_squeezing(self):
         for metric, kwargs, expected_score in zip(metrics, other_kwargs, expected_scores):
             metric_wrapper = MetricWrapper(
                 metric=metric,
+                multitask_handling="mean-per-label",
                 squeeze_targets=True,
                 target_to_int=True,
                 **kwargs,
diff --git a/tests/test_multitask_datamodule.py b/tests/test_multitask_datamodule.py
index 5cee26d61..ef66f10a2 100644
--- a/tests/test_multitask_datamodule.py
+++ b/tests/test_multitask_datamodule.py
@@ -1,11 +1,14 @@
 import unittest as ut
 from omegaconf import OmegaConf
-
+import pandas as pd
+import numpy as np
 import goli
 
 
 class Test_Multitask_DataModule(ut.TestCase):
-    def test_multitask_fromsmiles_dm(self):
+    def test_multitask_fromsmiles_dm(
+        self,
+    ):  # TODO: I think we can remove this as it tests tiny_zinc which only contain graph level labels
         """Cover similar testing as for the original data module."""
         df = goli.data.load_tiny_zinc()  # 100 molecules
 
@@ -28,6 +31,7 @@ def test_multitask_fromsmiles_dm(self):
         # Per-task arguments.
         dm_task_args_SA = {}
         dm_task_args_SA["df"] = df_micro_zinc_SA
+        dm_task_args_SA["task_level"] = "graph"
         dm_task_args_SA["smiles_col"] = "SMILES"
         dm_task_args_SA["label_cols"] = ["SA"]
         dm_task_args_SA["split_val"] = 0.2
@@ -41,6 +45,7 @@ def test_multitask_fromsmiles_dm(self):
 
         dm_task_args_logp = {}
         dm_task_args_logp["df"] = df_micro_zinc_logp
+        dm_task_args_logp["task_level"] = "graph"
         dm_task_args_logp["smiles_col"] = "SMILES"
         dm_task_args_logp["label_cols"] = ["logp"]
         dm_task_args_logp["split_val"] = 0.2
@@ -54,6 +59,7 @@ def test_multitask_fromsmiles_dm(self):
 
         dm_task_args_score = {}
         dm_task_args_score["df"] = df_micro_zinc_score
+        dm_task_args_score["task_level"] = "graph"
         dm_task_args_score["smiles_col"] = "SMILES"
         dm_task_args_score["label_cols"] = ["score"]
         dm_task_args_score["split_val"] = 0.2
@@ -109,9 +115,9 @@ def test_multitask_fromsmiles_dm(self):
             assert set(batch.keys()) == {"labels", "features"}
 
             # assert batch["labels"].shape == (16, 1)            # Single-task case
-            assert batch["labels"]["SA"].shape == (16, 1)
-            assert batch["labels"]["logp"].shape == (16, 1)
-            assert batch["labels"]["score"].shape == (16, 1)
+            assert batch["labels"]["graph_SA"].shape == (16, 1)
+            assert batch["labels"]["graph_logp"].shape == (16, 1)
+            assert batch["labels"]["graph_score"].shape == (16, 1)
 
     def test_multitask_fromsmiles_from_config(self):
         config = goli.load_config(name="zinc_default_multitask_pyg")
@@ -130,6 +136,14 @@ def test_multitask_fromsmiles_from_config(self):
         dm_args["task_specific_args"]["logp"]["df"] = df_micro_zinc_logp
         dm_args["task_specific_args"]["score"]["df"] = df_micro_zinc_score
 
+        dm_args["task_specific_args"]["SA"]["smiles_col"] = "SMILES"
+        dm_args["task_specific_args"]["logp"]["smiles_col"] = "SMILES"
+        dm_args["task_specific_args"]["score"]["smiles_col"] = "SMILES"
+
+        dm_args["task_specific_args"]["SA"]["label_cols"] = ["SA"]
+        dm_args["task_specific_args"]["logp"]["label_cols"] = ["logp"]
+        dm_args["task_specific_args"]["score"]["label_cols"] = ["score"]
+
         dm_args["task_specific_args"]["SA"]["df_path"] = None
         dm_args["task_specific_args"]["logp"]["df_path"] = None
         dm_args["task_specific_args"]["score"]["df_path"] = None
@@ -156,9 +170,9 @@ def test_multitask_fromsmiles_from_config(self):
             assert set(batch.keys()) == {"labels", "features"}
 
             # assert batch["labels"].shape == (16, 1)            # Single-task case
-            assert batch["labels"]["SA"].shape == (16, 1)
-            assert batch["labels"]["logp"].shape == (16, 1)
-            assert batch["labels"]["score"].shape == (16, 1)
+            assert batch["labels"]["graph_SA"].shape == (16, 1)
+            assert batch["labels"]["graph_logp"].shape == (16, 1)
+            assert batch["labels"]["graph_score"].shape == (16, 1)
 
     def test_multitask_fromsmiles_from_config_csv(self):
         config = goli.load_config(name="zinc_default_multitask_pyg")
@@ -183,9 +197,148 @@ def test_multitask_fromsmiles_from_config_csv(self):
             assert set(batch.keys()) == {"labels", "features"}
 
             # assert batch["labels"].shape == (16, 1)            # Single-task case
-            assert batch["labels"]["SA"].shape == (16, 1)
-            assert batch["labels"]["logp"].shape == (16, 1)
-            assert batch["labels"]["score"].shape == (16, 1)
+            assert batch["labels"]["graph_SA"].shape == (16, 1)
+            assert batch["labels"]["graph_logp"].shape == (16, 1)
+            assert batch["labels"]["graph_score"].shape == (16, 1)
+
+    def test_multitask_fromsmiles_from_config_parquet(self):
+        config = goli.load_config(name="fake_multilevel_multitask_pyg")
+
+        dm_args = OmegaConf.to_container(config.datamodule.args, resolve=True)
+        dm = goli.data.MultitaskFromSmilesDataModule(**dm_args)
+
+        dm.prepare_data()
+        dm.setup()
+
+        # self.assertEqual(len(dm), 100)                      # Should this have a fixed value for when it's initialized? MTL dataset only gets created after.
+        self.assertEqual(len(dm.train_ds), 602)  # type: ignore
+        self.assertEqual(len(dm.val_ds), 201)  # type: ignore
+        self.assertEqual(len(dm.test_ds), 201)  # type: ignore
+        # assert dm.num_node_feats == 50
+        # assert dm.num_edge_feats == 6
+
+        for dl in [dm.train_dataloader(), dm.val_dataloader(), dm.test_dataloader()]:
+            dl = dm.train_dataloader()
+            it = iter(dl)
+            batch = next(it)
+
+            assert set(batch.keys()) == {"labels", "features"}
+
+            # assert batch["labels"].shape == (16, 1)            # Single-task case
+            assert batch["labels"]["graph_SA"].shape == (16, 1)
+            assert batch["labels"]["node_logp"].shape == (
+                batch["features"].feat.size(0),
+                2,
+            )  # test node level
+            assert batch["labels"]["edge_score"].shape == (
+                batch["features"].edge_feat.size(0),
+                2,
+            )  # test edge level
+
+    def test_multitask_with_missing_fromsmiles_from_config_parquet(self):
+        config = goli.load_config(name="fake_and_missing_multilevel_multitask_pyg")
+
+        dm_args = OmegaConf.to_container(config.datamodule.args, resolve=True)
+        dm = goli.data.MultitaskFromSmilesDataModule(**dm_args)
+
+        dm.prepare_data()
+        dm.setup()
+
+        # self.assertEqual(len(dm), 100)                      # Should this have a fixed value for when it's initialized? MTL dataset only gets created after.
+        self.assertEqual(len(dm.train_ds), 602)  # type: ignore
+        self.assertEqual(len(dm.val_ds), 201)  # type: ignore
+        self.assertEqual(len(dm.test_ds), 201)  # type: ignore
+        # assert dm.num_node_feats == 50
+        # assert dm.num_edge_feats == 6
+
+        for dl in [dm.train_dataloader(), dm.val_dataloader(), dm.test_dataloader()]:
+            dl = dm.train_dataloader()
+            it = iter(dl)
+            batch = next(it)
+
+            assert set(batch.keys()) == {"labels", "features"}
+
+            # assert batch["labels"].shape == (16, 1)            # Single-task case
+            assert batch["labels"]["graph_SA"].shape == (16, 1)
+            assert batch["labels"]["node_logp"].shape == (
+                batch["features"].feat.size(0),
+                2,
+            )  # test node level
+            assert batch["labels"]["edge_score"].shape == (
+                batch["features"].edge_feat.size(0),
+                2,
+            )  # test edge level
+
+    def test_extract_graph_level_singletask(self):
+        df = pd.read_parquet(f"tests/converted_fake_multilevel_data.parquet")
+        num_graphs = len(df)
+        label_cols = ["graph_label"]
+        output = goli.data.datamodule.extract_labels(df, "graph", label_cols)
+
+        assert isinstance(output, np.ndarray)
+        assert len(output.shape) == 2
+        assert output.shape[0] == num_graphs
+        assert output.shape[1] == 1
+
+    def test_extract_graph_level_multitask(self):
+        df = pd.read_parquet(f"tests/converted_fake_multilevel_data.parquet")
+        num_graphs = len(df)
+        label_cols = ["graph_label", "graph_label"]
+        output = goli.data.datamodule.extract_labels(df, "graph", label_cols)
+
+        assert isinstance(output, np.ndarray)
+        assert len(output.shape) == 2
+        assert output.shape[0] == num_graphs
+        assert output.shape[1] == len(label_cols)
+
+    def test_extract_graph_level_multitask_missing_cols(self):
+        df = pd.read_parquet(f"tests/converted_fake_multilevel_data.parquet")
+        num_graphs = len(df)
+        label_cols = ["graph_label", "graph_label"]
+
+        drop_index = [2, 5, 21, 237, 192, 23, 127, 11]
+        for replace in [1, 2]:
+            for missing_col in label_cols[:replace]:
+                df[missing_col].iloc[drop_index] = None
+
+            output = goli.data.datamodule.extract_labels(df, "graph", label_cols)
+
+            assert isinstance(output, np.ndarray)
+            assert len(output.shape) == 2
+            assert output.shape[0] == num_graphs
+            assert output.shape[1] == len(label_cols)
+
+    def test_non_graph_level_extract_labels(self):
+        df = pd.read_parquet(f"tests/converted_fake_multilevel_data.parquet")
+
+        for level in ["node", "edge", "nodepair"]:
+            label_cols = [f"{level}_label_{suffix}" for suffix in ["list", "np"]]
+            output = goli.data.datamodule.extract_labels(df, level, label_cols)
+
+            assert isinstance(output, list)
+            assert len(output[0].shape) == 2
+            assert output[0].shape[1] == len(label_cols)
+
+    def test_non_graph_level_extract_labels_missing_cols(self):
+        df = pd.read_parquet(f"tests/converted_fake_multilevel_data.parquet")
+
+        for level in ["node", "edge", "nodepair"]:
+            label_cols = [f"{level}_label_{suffix}" for suffix in ["list", "np"]]
+            drop_index = [2, 5, 21, 237, 192, 23, 127, 11]
+            for replace in [1, 2]:
+                for missing_col in label_cols[:replace]:
+                    df.loc[drop_index, missing_col] = None
+
+                output = goli.data.datamodule.extract_labels(df, level, label_cols)
+
+                for idx in drop_index:
+                    assert len(output[idx].shape) == 2
+                    assert output[idx].shape[1] == len(label_cols)
+
+                    # Check that number of labels is adjusted correctly
+                    if replace == 1:
+                        non_missing_col = label_cols[1]
+                        assert output[idx].shape[0] == len(df[non_missing_col][idx])
 
 
 if __name__ == "__main__":