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

Issue:
Deal with NaN featurization
Deal with molecules that fail a part of the featurization, such as `[He][He]`.

Add a parameter `nan_mask` to the featurization of both nodes and edges with different possible attributes:

- `"raise"`: DEFAULT. Raise an error when there is a nan in the featurization
- `"warn"`: Raise a warning when there is a nan in the featurization
- `None`: Don't do anything
- Floating value: Replace nans by the specified value
Save a compressed dataset file
https://github.com/valence-discovery/goli/blob/cb3bbce8c92967bd5e47f531a5eb5a0f337e6379/goli/data/datamodule.py#L535

Some large datasets require >30GB of storage, it's way too much, and too long and large to download from an online file. For example, the processed `htsfp-pcba` dataset takes about 15GB, but compressed it takes 860MB, a very significant difference, both in storage and download time.

Allow the saving and loading to use compressed format such as `.cache.gz`
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:

Fix patch:
diff --git a/docs/tutorials/basics/implementing_gnn_layers.ipynb b/docs/tutorials/basics/implementing_gnn_layers.ipynb
index ec154a81a..d2d585dab 100644
--- a/docs/tutorials/basics/implementing_gnn_layers.ipynb
+++ b/docs/tutorials/basics/implementing_gnn_layers.ipynb
@@ -10,13 +10,15 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.8-final"
+   "version": "3.8.8"
   },
   "orig_nbformat": 2,
   "kernelspec": {
    "name": "python3",
-   "display_name": "Python 3",
-   "language": "python"
+   "display_name": "Python 3.8.8 64-bit ('goli': conda)"
+  },
+  "interpreter": {
+   "hash": "f4a99d018a205fcbcc0480c84566beaebcb91b08d0414b39a842df533e2a1d25"
   }
  },
  "nbformat": 4,
@@ -26,7 +28,9 @@
    "source": [
     "# Creating GNN layers\n",
     "\n",
-    "In this example, you will learn how to create a custom GNN class such that it is compatible with the flexible network architecture `goli.nn.architectures.FeedForwardDGL`. We will first start by a simple layer that does not use edges, to a more complex layer that uses edges.\n",
+    "One of the primary advantage of the current library is the fact that GNN layers are independant from model architecture, thus allowing more flexibility with the code by easily swapping different layer types as hyper-parameters. However, this requires that the layers be implemented using the DGL library, and must be inherited from the class `BaseDGLLayer`, which standardizes the inputs, outputs and properties of the layers. Thus, the architecture can be handled independantly using the class `FeedForwardDGL`, or any similar custom class.\n",
+    "\n",
+    "We will first start by a simple layer that does not use edges, to a more complex layer that uses edges.\n",
     "\n",
     "Since these examples are built on top of DGL, we recommend looking at their [library](https://docs.dgl.ai/en/0.5.x/index.html) for more info. "
    ],
@@ -115,7 +119,7 @@
     "\n",
     "Here, we will show how to create a GNN layer that does a mean aggregation on the neighbouring features.\n",
     "\n",
-    "First, for the layer to be fully compatible with the flexible architecture provided by `goli.nn.architectures.FeedForwardDGL`, it needs to inhering from the class `goli.nn.dgl_layers.BaseDGLLayer`. This base-layer has multiple virtual methods that must be implemented in any class that inherits from it.\n",
+    "First, for the layer to be fully compatible with the flexible architecture provided by `FeedForwardDGL`, it needs to inherit from the class `BaseDGLLayer`. This base-layer has multiple virtual methods that must be implemented in any class that inherits from it.\n",
     "\n",
     "The virtual methods are below\n",
     "\n",
diff --git a/docs/tutorials/basics/making_gnn_networks.ipynb b/docs/tutorials/basics/making_gnn_networks.ipynb
index 5bff7ffd7..0c33e07b8 100644
--- a/docs/tutorials/basics/making_gnn_networks.ipynb
+++ b/docs/tutorials/basics/making_gnn_networks.ipynb
@@ -6,7 +6,11 @@
    "source": [
     "# Making GNN Networks\n",
     "\n",
-    "In this example, you will learn how to easily build a full GNN network using any kind of GNN layer."
+    "In this example, you will learn how to easily build a full GNN network using any kind of GNN layer. This tutorial uses the architecture defined by the class `FullDGLNetwork`, which is compatible with any layer that inherits from `BaseDGLLayer`.\n",
+    "\n",
+    "`FullDGLNetwork` is an architecture that takes as input node features and (optionally) edge features. It applies a pre-MLP on both sets of features, then it passes it into a main GNN network, and finally applies a post-MLP to produce the final output (either node predictions or graph property predictions).\n",
+    "\n",
+    "The network is very easy to built via a dictionnary of parameter that allow to custom each part of the network."
    ]
   },
   {
@@ -40,7 +44,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "We will first create some simple batched graphs that will be used accross the examples."
+    "We will first create some simple batched graphs that will be used accross the examples. Here, `bg` is a batch containing 2 graphs with random node features."
    ]
   },
   {
@@ -157,7 +161,9 @@
     "\n",
     "Once the network is defined, we only need to run the forward pass on the input graphs to get a prediction.\n",
     "\n",
-    "The network will handle the node and edge features depending on it's parameters and layer type."
+    "The network will handle the node and edge features depending on it's parameters and layer type.\n",
+    "\n",
+    "Chosing between graph property prediction and node property prediction depends on the parameter given by `gnn_kwargs[\"pooling\"]`."
    ]
   },
   {
@@ -196,12 +202,7 @@
  "metadata": {
   "kernelspec": {
    "name": "python3",
-   "display_name": "Python 3.8.8 64-bit ('goli': conda)",
-   "metadata": {
-    "interpreter": {
-     "hash": "0db34acba4cc11eb0fa2be1290630bee7d7b89811a824f1a15cd26acd4633567"
-    }
-   }
+   "display_name": "Python 3.8.8 64-bit ('goli': conda)"
   },
   "language_info": {
    "codemirror_mode": {
@@ -213,7 +214,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.8-final"
+   "version": "3.8.8"
   },
   "widgets": {
    "application/vnd.jupyter.widget-state+json": {
@@ -221,6 +222,9 @@
     "version_major": 2,
     "version_minor": 0
    }
+  },
+  "interpreter": {
+   "hash": "f4a99d018a205fcbcc0480c84566beaebcb91b08d0414b39a842df533e2a1d25"
   }
  },
  "nbformat": 4,
diff --git a/docs/tutorials/basics/using_gnn_layers.ipynb b/docs/tutorials/basics/using_gnn_layers.ipynb
index 31aedcf0d..345c7cbfc 100644
--- a/docs/tutorials/basics/using_gnn_layers.ipynb
+++ b/docs/tutorials/basics/using_gnn_layers.ipynb
@@ -10,13 +10,15 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.8-final"
+   "version": "3.8.8"
   },
   "orig_nbformat": 2,
   "kernelspec": {
    "name": "python3",
-   "display_name": "Python 3",
-   "language": "python"
+   "display_name": "Python 3.8.8 64-bit ('goli': conda)"
+  },
+  "interpreter": {
+   "hash": "f4a99d018a205fcbcc0480c84566beaebcb91b08d0414b39a842df533e2a1d25"
   }
  },
  "nbformat": 4,
@@ -26,7 +28,9 @@
    "source": [
     "# Using GNN layers\n",
     "\n",
-    "In this example, you will learn how to use the GCN, GIN, Gated-GCN and PNA layers in a simple `forward` context."
+    "The current library implements multiple state-of-the-art graph neural networks. In this tutorial, you will learn how to use the **GCN**, **GIN**, **Gated-GCN** and **PNA** layers in a simple `forward` context.\n",
+    "\n",
+    "Other layers such as **DGN** require additional positional encoding to work."
    ],
    "cell_type": "markdown",
    "metadata": {}
@@ -66,7 +70,7 @@
   },
   {
    "source": [
-    "We will first create some simple batched graphs that will be used accross the examples."
+    "We will first create some simple batched graphs that will be used accross the examples. Here, `bg` is a batch of 2 graphs containing random node features `bg.ndata[\"h\"]` and edge features `bg.edata[\"e\"]`."
    ],
    "cell_type": "markdown",
    "metadata": {}
@@ -287,7 +291,7 @@
     "\n",
     "## PNA-conv\n",
     "\n",
-    "First, let's focus on the PNA-conv. In this case, it works exactly as the GCN and GIN methods."
+    "First, let's focus on the PNA-conv. In this case, it works exactly as the GCN and GIN methods. Although not presented in the example, PNA-conv also supports edge features by concatenating them to the node features during convolution."
    ],
    "cell_type": "markdown",
    "metadata": {}
@@ -327,7 +331,7 @@
    "source": [
     "## PNA-msgpass\n",
     "\n",
-    "The PNA message passing is typically more powerful that the convolutional one, and it supports edges as inputs, but doesn't output edges. It's usage is very similar to the *PNA-conv*, but we need to specify the edge dimensions and features."
+    "The PNA message passing is typically more powerful that the convolutional one, and it supports edges as inputs, but doesn't output edges. It's usage is very similar to the *PNA-conv*. Here, we also present the option to specify the edge dimensions and features."
    ],
    "cell_type": "markdown",
    "metadata": {}
diff --git a/docs/tutorials/model_training/simple-molecular-model.ipynb b/docs/tutorials/model_training/simple-molecular-model.ipynb
index c383499b7..160e8cd05 100644
--- a/docs/tutorials/model_training/simple-molecular-model.ipynb
+++ b/docs/tutorials/model_training/simple-molecular-model.ipynb
@@ -8,6 +8,8 @@
     "\n",
     "This tutorial will walk you through how to use a configuration file to define all the parameters of a model and of the trainer. This tutorial focuses on training from SMILES data in a CSV format.\n",
     "\n",
+    "There are multiple examples of YAML files located in the folder `goli/expts` that one can refer to when training a new model. The file `config_ZINC_bench_gnn.yaml` shows an example of single task regression from a CSV file provided by goli. And the file `config_molpcba.yaml` shows an example of a multi-task classification on a dataset provided by OGB with some missing data.\n",
+    "\n",
     "## Creating the yaml file\n",
     "\n",
     "The first step is to create a YAML file containing all the required configurations, with an example given at `goli/expts/config_micro_ZINC.yaml`. We will go through each part of the configurations."
@@ -40,8 +42,8 @@
    "metadata": {},
    "outputs": [
     {
-     "name": "stdout",
      "output_type": "stream",
+     "name": "stdout",
      "text": [
       "Yaml file loaded\n"
      ]
@@ -70,13 +72,10 @@
    "metadata": {},
    "outputs": [
     {
-     "name": "stdout",
      "output_type": "stream",
+     "name": "stdout",
      "text": [
-      "constants:\n",
-      "  seed: 42\n",
-      "  raise_train_error: true\n",
-      "\n"
+      "constants:\n  seed: 42\n  raise_train_error: true\n\n"
      ]
     }
    ],
@@ -101,38 +100,10 @@
    "metadata": {},
    "outputs": [
     {
-     "name": "stdout",
      "output_type": "stream",
+     "name": "stdout",
      "text": [
-      "datamodule:\n",
-      "  df_path: goli/data/micro_ZINC/micro_ZINC.csv\n",
-      "  cache_data_path: goli/data/cache/micro_ZINC/full.cache\n",
-      "  label_cols:\n",
-      "  - score\n",
-      "  smiles_col: SMILES\n",
-      "  featurization_n_jobs: -1\n",
-      "  featurization_progress: true\n",
-      "  featurization:\n",
-      "    atom_property_list_onehot:\n",
-      "    - atomic-number\n",
-      "    - valence\n",
-      "    atom_property_list_float:\n",
-      "    - mass\n",
-      "    - electronegativity\n",
-      "    - in-ring\n",
-      "    edge_property_list: []\n",
-      "    add_self_loop: false\n",
-      "    explicit_H: false\n",
-      "    use_bonds_weights: false\n",
-      "  split_val: 0.2\n",
-      "  split_test: 0.2\n",
-      "  split_seed: 42\n",
-      "  splits_path: null\n",
-      "  batch_size_train_val: 128\n",
-      "  batch_size_test: 256\n",
-      "  num_workers: -1\n",
-      "  pin_memory: false\n",
-      "\n"
+      "datamodule:\n  module_type: DGLFromSmilesDataModule\n  args:\n    df_path: goli/data/micro_ZINC/micro_ZINC.csv\n    cache_data_path: goli/data/cache/micro_ZINC/full.cache.gz\n    label_cols:\n    - score\n    smiles_col: SMILES\n    featurization_n_jobs: -1\n    featurization_progress: true\n    featurization:\n      atom_property_list_onehot:\n      - atomic-number\n      - valence\n      atom_property_list_float:\n      - mass\n      - electronegativity\n      - in-ring\n      edge_property_list:\n      - bond-type-onehot\n      - stereo\n      - in-ring\n      add_self_loop: false\n      explicit_H: false\n      use_bonds_weights: false\n      pos_encoding_as_features:\n        pos_type: laplacian_eigvec\n        num_pos: 3\n        normalization: none\n        disconnected_comp: true\n      pos_encoding_as_directions:\n        pos_type: laplacian_eigvec\n        num_pos: 3\n        normalization: none\n        disconnected_comp: true\n    split_val: 0.2\n    split_test: 0.2\n    split_seed: 42\n    splits_path: null\n    batch_size_train_val: 128\n    batch_size_test: 128\n    num_workers: 0\n    pin_memory: false\n    persistent_workers: false\n\n"
      ]
     }
    ],
@@ -163,55 +134,10 @@
    "metadata": {},
    "outputs": [
     {
-     "name": "stdout",
      "output_type": "stream",
+     "name": "stdout",
      "text": [
-      "architecture:\n",
-      "  model_type: fulldglnetwork\n",
-      "  pre_nn:\n",
-      "    out_dim: 32\n",
-      "    hidden_dims: 32\n",
-      "    depth: 1\n",
-      "    activation: relu\n",
-      "    last_activation: none\n",
-      "    dropout: 0.1\n",
-      "    batch_norm: true\n",
-      "    last_batch_norm: true\n",
-      "    residual_type: none\n",
-      "  gnn:\n",
-      "    out_dim: 32\n",
-      "    hidden_dims: 32\n",
-      "    depth: 4\n",
-      "    activation: relu\n",
-      "    last_activation: none\n",
-      "    dropout: 0.1\n",
-      "    batch_norm: true\n",
-      "    last_batch_norm: true\n",
-      "    residual_type: simple\n",
-      "    pooling: sum\n",
-      "    virtual_node: sum\n",
-      "    layer_type: pna-msgpass\n",
-      "    layer_kwargs:\n",
-      "      aggregators:\n",
-      "      - mean\n",
-      "      - max\n",
-      "      - min\n",
-      "      - std\n",
-      "      scalers:\n",
-      "      - identity\n",
-      "      - amplification\n",
-      "      - attenuation\n",
-      "  post_nn:\n",
-      "    out_dim: 1\n",
-      "    hidden_dims: 32\n",
-      "    depth: 2\n",
-      "    activation: relu\n",
-      "    last_activation: none\n",
-      "    dropout: 0.1\n",
-      "    batch_norm: true\n",
-      "    last_batch_norm: false\n",
-      "    residual_type: none\n",
-      "\n"
+      "architecture:\n  model_type: fulldglnetwork\n  pre_nn:\n    out_dim: 32\n    hidden_dims: 32\n    depth: 1\n    activation: relu\n    last_activation: none\n    dropout: 0.1\n    batch_norm: true\n    last_batch_norm: true\n    residual_type: none\n  pre_nn_edges:\n    out_dim: 16\n    hidden_dims: 16\n    depth: 2\n    activation: relu\n    last_activation: none\n    dropout: 0.1\n    batch_norm: true\n    last_batch_norm: true\n    residual_type: none\n  gnn:\n    out_dim: 32\n    hidden_dims: 32\n    depth: 4\n    activation: relu\n    last_activation: none\n    dropout: 0.1\n    batch_norm: true\n    last_batch_norm: true\n    residual_type: simple\n    pooling:\n    - sum\n    - max\n    - dir1\n    virtual_node: sum\n    layer_type: dgn-msgpass\n    layer_kwargs:\n      aggregators:\n      - mean\n      - max\n      - dir1/dx_abs\n      - dir1/smooth\n      scalers:\n      - identity\n      - amplification\n      - attenuation\n  post_nn:\n    out_dim: 1\n    hidden_dims: 32\n    depth: 2\n    activation: relu\n    last_activation: none\n    dropout: 0.1\n    batch_norm: true\n    last_batch_norm: false\n    residual_type: none\n\n"
      ]
     }
    ],
@@ -234,28 +160,10 @@
    "metadata": {},
    "outputs": [
     {
-     "name": "stdout",
      "output_type": "stream",
+     "name": "stdout",
      "text": [
-      "predictor:\n",
-      "  metrics_on_progress_bar:\n",
-      "  - mae\n",
-      "  - pearsonr\n",
-      "  - f1 > 3\n",
-      "  - precision > 3\n",
-      "  loss_fun: mse\n",
-      "  random_seed: 42\n",
-      "  optim_kwargs:\n",
-      "    lr: 0.01\n",
-      "    weight_decay: 1.0e-07\n",
-      "  lr_reduce_on_plateau_kwargs:\n",
-      "    factor: 0.5\n",
-      "    patience: 7\n",
-      "  scheduler_kwargs:\n",
-      "    monitor: loss/val\n",
-      "    frequency: 1\n",
-      "  target_nan_mask: 0\n",
-      "\n"
+      "predictor:\n  metrics_on_progress_bar:\n  - mae\n  - pearsonr\n  - f1 > 3\n  - precision > 3\n  loss_fun: mse\n  random_seed: 42\n  optim_kwargs:\n    lr: 0.01\n    weight_decay: 1.0e-07\n  lr_reduce_on_plateau_kwargs:\n    factor: 0.5\n    patience: 7\n  scheduler_kwargs:\n    monitor: loss/val\n    frequency: 1\n  target_nan_mask: 0\n\n"
      ]
     }
    ],
@@ -282,43 +190,10 @@
    "metadata": {},
    "outputs": [
     {
-     "name": "stdout",
      "output_type": "stream",
+     "name": "stdout",
      "text": [
-      "metrics:\n",
-      "- name: mae\n",
-      "  metric: mae\n",
-      "  threshold_kwargs: null\n",
-      "- name: pearsonr\n",
-      "  metric: pearsonr\n",
-      "  threshold_kwargs: null\n",
-      "- name: f1 > 3\n",
-      "  metric: f1\n",
-      "  num_classes: 2\n",
-      "  average: micro\n",
-      "  threshold_kwargs:\n",
-      "    operator: greater\n",
-      "    threshold: 3\n",
-      "    th_on_preds: true\n",
-      "    th_on_target: true\n",
-      "- name: f1 > 5\n",
-      "  metric: f1\n",
-      "  num_classes: 2\n",
-      "  average: micro\n",
-      "  threshold_kwargs:\n",
-      "    operator: greater\n",
-      "    threshold: 5\n",
-      "    th_on_preds: true\n",
-      "    th_on_target: true\n",
-      "- name: precision > 3\n",
-      "  metric: precision\n",
-      "  class_reduction: micro\n",
-      "  threshold_kwargs:\n",
-      "    operator: greater\n",
-      "    threshold: 3\n",
-      "    th_on_preds: true\n",
-      "    th_on_target: true\n",
-      "\n"
+      "metrics:\n- name: mae\n  metric: mae\n  threshold_kwargs: null\n- name: pearsonr\n  metric: pearsonr\n  threshold_kwargs: null\n- name: f1 > 3\n  metric: f1\n  num_classes: 2\n  average: micro\n  threshold_kwargs:\n    operator: greater\n    threshold: 3\n    th_on_preds: true\n    th_on_target: true\n    target_to_int: true\n- name: f1 > 5\n  metric: f1\n  num_classes: 2\n  average: micro\n  threshold_kwargs:\n    operator: greater\n    threshold: 5\n    th_on_preds: true\n    th_on_target: true\n    target_to_int: true\n- name: precision > 3\n  metric: precision\n  class_reduction: micro\n  threshold_kwargs:\n    operator: greater\n    threshold: 3\n    th_on_preds: true\n    th_on_target: true\n    target_to_int: true\n\n"
      ]
     }
    ],
@@ -341,29 +216,10 @@
    "metadata": {},
    "outputs": [
     {
-     "name": "stdout",
      "output_type": "stream",
+     "name": "stdout",
      "text": [
-      "trainer:\n",
-      "  logger:\n",
-      "    save_dir: logs/micro_ZINC\n",
-      "  early_stopping:\n",
-      "    monitor: loss/val\n",
-      "    min_delta: 0\n",
-      "    patience: 10\n",
-      "    mode: min\n",
-      "  model_checkpoint:\n",
-      "    dirpath: models_checkpoints/micro_ZINC/\n",
-      "    filename: bob\n",
-      "    monitor: loss/val\n",
-      "    mode: min\n",
-      "    save_top_k: 1\n",
-      "    period: 1\n",
-      "  trainer:\n",
-      "    max_epochs: 25\n",
-      "    min_epochs: 5\n",
-      "    gpus: 1\n",
-      "\n"
+      "trainer:\n  logger:\n    save_dir: logs/micro_ZINC\n  early_stopping:\n    monitor: loss/val\n    min_delta: 0\n    patience: 10\n    mode: min\n  model_checkpoint:\n    dirpath: models_checkpoints/micro_ZINC/\n    filename: model\n    monitor: loss/val\n    mode: min\n    save_top_k: 1\n    period: 1\n  trainer:\n    max_epochs: 25\n    min_epochs: 5\n    gpus: 1\n\n"
      ]
     }
    ],
@@ -386,24 +242,33 @@
    "metadata": {},
    "outputs": [
     {
-     "name": "stderr",
      "output_type": "stream",
+     "name": "stderr",
      "text": [
       "Using backend: pytorch\n",
-      "2021-03-25 09:44:37.314 | WARNING  | goli.config._loader:load_trainer:111 - Number of GPUs selected is `1`, but will be ignored since no GPU are available on this device\n",
-      "/home/dominique/anaconda3/envs/goli/lib/python3.8/site-packages/pytorch_lightning/utilities/distributed.py:51: UserWarning: Checkpoint directory models_checkpoints/micro_ZINC/ exists and is not empty.\n",
+      "c:\\users\\domin\\documents\\gits\\goli_windows\\goli\\features\\spectral.py:43: ComplexWarning: Casting complex values to real discards the imaginary part\n",
+      "  eigvecs[comp, :] = this_eigvecs\n",
+      "c:\\users\\domin\\documents\\gits\\goli_windows\\goli\\features\\spectral.py:44: ComplexWarning: Casting complex values to real discards the imaginary part\n",
+      "  eigvals_tile[comp, :] = this_eigvals\n",
+      "2021-06-21 12:39:11.913 | WARNING  | goli.config._loader:load_trainer:126 - Number of GPUs selected is `1`, but will be ignored since no GPU are available on this device\n",
+      "C:\\Users\\Domin\\miniconda3\\envs\\goli\\lib\\site-packages\\pytorch_lightning\\utilities\\distributed.py:68: UserWarning: Checkpoint directory models_checkpoints/micro_ZINC/ exists and is not empty.\n",
       "  warnings.warn(*args, **kwargs)\n",
       "GPU available: False, used: False\n",
-      "TPU available: None, using: 0 TPU cores\n",
-      "2021-03-25 09:44:37.331 | INFO     | goli.data.datamodule:prepare_data:153 - Reload data from goli/data/cache/micro_ZINC/full.cache.\n",
+      "TPU available: False, using: 0 TPU cores\n",
+      "2021-06-21 12:39:11.927 | INFO     | goli.data.datamodule:_load_from_cache:567 - Try reloading the data module from goli/data/cache/micro_ZINC/full.cache.gz.\n",
       "\n",
       "datamodule:\n",
       " name: DGLFromSmilesDataModule\n",
       "len: 1000\n",
+      "train_size: null\n",
+      "val_size: null\n",
+      "test_size: null\n",
       "batch_size_train_val: 128\n",
-      "batch_size_test: 256\n",
+      "batch_size_test: 128\n",
       "num_node_feats: 55\n",
-      "num_edge_feats: 0\n",
+      "num_node_feats_with_positional_encoding: 58\n",
+      "num_edge_feats: 13\n",
+      "num_labels: 1\n",
       "collate_fn: goli_collate_fn\n",
       "featurization:\n",
       "  atom_property_list_onehot:\n",
@@ -413,469 +278,236 @@
       "  - mass\n",
       "  - electronegativity\n",
       "  - in-ring\n",
-      "  edge_property_list: []\n",
+      "  edge_property_list:\n",
+      "  - bond-type-onehot\n",
+      "  - stereo\n",
+      "  - in-ring\n",
       "  add_self_loop: false\n",
       "  explicit_H: false\n",
       "  use_bonds_weights: false\n",
+      "  pos_encoding_as_features:\n",
+      "    pos_type: laplacian_eigvec\n",
+      "    num_pos: 3\n",
+      "    normalization: none\n",
+      "    disconnected_comp: true\n",
+      "  pos_encoding_as_directions:\n",
+      "    pos_type: laplacian_eigvec\n",
+      "    num_pos: 3\n",
+      "    normalization: none\n",
+      "    disconnected_comp: true\n",
       " \n",
       "\n",
       "{'mae': mean_absolute_error, 'pearsonr': pearsonr, 'f1 > 3': f1(>3), 'f1 > 5': f1(>5), 'precision > 3': precision(>3)}\n",
       "DGL_GNN\n",
       "---------\n",
       "    pre-NN(depth=1, ResidualConnectionNone)\n",
-      "        [FCLayer[55 -> 32]\n",
+      "        [FCLayer[58 -> 32]\n",
+      "    \n",
+      "    pre-NN-edges(depth=2, ResidualConnectionNone)\n",
+      "        [FCLayer[13 -> 16 -> 16]\n",
       "    \n",
       "    GNN(depth=4, ResidualConnectionSimple(skip_steps=1))\n",
-      "        PNAMessagePassingLayer[32 -> 32 -> 32 -> 32 -> 32]\n",
-      "        -> Pooling(sum) -> FCLayer(32 -> 32, activation=None)\n",
+      "        DGNMessagePassingLayer[32 -> 32 -> 32 -> 32 -> 32]\n",
+      "        -> Pooling(['sum', 'max', 'dir1']) -> FCLayer(96 -> 32, activation=None)\n",
       "    \n",
       "    post-NN(depth=2, ResidualConnectionNone)\n",
       "        [FCLayer[32 -> 32 -> 1]\n",
-      "   | Name                            | Type                     | Params\n",
-      "------------------------------------------------------------------------------\n",
-      "0  | model                           | FullDGLNetwork           | 69.7 K\n",
-      "1  | model.pre_nn                    | FeedForwardNN            | 1.9 K \n",
-      "2  | model.pre_nn.activation         | ReLU                     | 0     \n",
-      "3  | model.pre_nn.residual_layer     | ResidualConnectionNone   | 0     \n",
-      "4  | model.pre_nn.layers             | ModuleList               | 1.9 K \n",
-      "5  | model.pre_nn.layers.0           | FCLayer                  | 1.9 K \n",
-      "6  | model.gnn                       | FeedForwardDGL           | 66.7 K\n",
-      "7  | model.gnn.activation            | ReLU                     | 0     \n",
-      "8  | model.gnn.layers                | ModuleList               | 62.2 K\n",
-      "9  | model.gnn.layers.0              | PNAMessagePassingLayer   | 15.6 K\n",
-      "10 | model.gnn.layers.1              | PNAMessagePassingLayer   | 15.6 K\n",
-      "11 | model.gnn.layers.2              | PNAMessagePassingLayer   | 15.6 K\n",
-      "12 | model.gnn.layers.3              | PNAMessagePassingLayer   | 15.6 K\n",
-      "13 | model.gnn.virtual_node_layers   | ModuleList               | 3.4 K \n",
-      "14 | model.gnn.virtual_node_layers.0 | VirtualNode              | 1.1 K \n",
-      "15 | model.gnn.virtual_node_layers.1 | VirtualNode              | 1.1 K \n",
-      "16 | model.gnn.virtual_node_layers.2 | VirtualNode              | 1.1 K \n",
-      "17 | model.gnn.residual_layer        | ResidualConnectionSimple | 0     \n",
-      "18 | model.gnn.global_pool_layer     | ModuleListConcat         | 0     \n",
-      "19 | model.gnn.global_pool_layer.0   | SumPooling               | 0     \n",
-      "20 | model.gnn.out_linear            | FCLayer                  | 1.1 K \n",
-      "21 | model.gnn.out_linear.linear     | Linear                   | 1.1 K \n",
-      "22 | model.gnn.out_linear.dropout    | Dropout                  | 0     \n",
-      "23 | model.gnn.out_linear.batch_norm | BatchNorm1d              | 64    \n",
-      "24 | model.post_nn                   | FeedForwardNN            | 1.2 K \n",
-      "25 | model.post_nn.activation        | ReLU                     | 0     \n",
-      "26 | model.post_nn.residual_layer    | ResidualConnectionNone   | 0     \n",
-      "27 | model.post_nn.layers            | ModuleList               | 1.2 K \n",
-      "28 | model.post_nn.layers.0          | FCLayer                  | 1.1 K \n",
-      "29 | model.post_nn.layers.1          | FCLayer                  | 33    \n",
-      "30 | loss_fun                        | MSELoss                  | 0     \n",
-      "------------------------------------------------------------------------------\n",
-      "69.7 K    Trainable params\n",
+      "   | Name                              | Type                     | Params\n",
+      "--------------------------------------------------------------------------------\n",
+      "0  | model                             | FullDGLNetwork           | 74.4 K\n",
+      "1  | model.pre_nn                      | FeedForwardNN            | 2.0 K \n",
+      "2  | model.pre_nn.activation           | ReLU                     | 0     \n",
+      "3  | model.pre_nn.residual_layer       | ResidualConnectionNone   | 0     \n",
+      "4  | model.pre_nn.layers               | ModuleList               | 2.0 K \n",
+      "5  | model.pre_nn.layers.0             | FCLayer                  | 2.0 K \n",
+      "6  | model.pre_nn_edges                | FeedForwardNN            | 560   \n",
+      "7  | model.pre_nn_edges.activation     | ReLU                     | 0     \n",
+      "8  | model.pre_nn_edges.residual_layer | ResidualConnectionNone   | 0     \n",
+      "9  | model.pre_nn_edges.layers         | ModuleList               | 560   \n",
+      "10 | model.pre_nn_edges.layers.0       | FCLayer                  | 256   \n",
+      "11 | model.pre_nn_edges.layers.1       | FCLayer                  | 304   \n",
+      "12 | model.gnn                         | FeedForwardDGL           | 70.8 K\n",
+      "13 | model.gnn.activation              | ReLU                     | 0     \n",
+      "14 | model.gnn.layers                  | ModuleList               | 64.3 K\n",
+      "15 | model.gnn.layers.0                | DGNMessagePassingLayer   | 16.1 K\n",
+      "16 | model.gnn.layers.1                | DGNMessagePassingLayer   | 16.1 K\n",
+      "17 | model.gnn.layers.2                | DGNMessagePassingLayer   | 16.1 K\n",
+      "18 | model.gnn.layers.3                | DGNMessagePassingLayer   | 16.1 K\n",
+      "19 | model.gnn.virtual_node_layers     | ModuleList               | 3.4 K \n",
+      "20 | model.gnn.virtual_node_layers.0   | VirtualNode              | 1.1 K \n",
+      "21 | model.gnn.virtual_node_layers.1   | VirtualNode              | 1.1 K \n",
+      "22 | model.gnn.virtual_node_layers.2   | VirtualNode              | 1.1 K \n",
+      "23 | model.gnn.residual_layer          | ResidualConnectionSimple | 0     \n",
+      "24 | model.gnn.global_pool_layer       | ModuleListConcat         | 0     \n",
+      "25 | model.gnn.global_pool_layer.0     | SumPooling               | 0     \n",
+      "26 | model.gnn.global_pool_layer.1     | MaxPooling               | 0     \n",
+      "27 | model.gnn.global_pool_layer.2     | DirPooling               | 0     \n",
+      "28 | model.gnn.out_linear              | FCLayer                  | 3.2 K \n",
+      "29 | model.gnn.out_linear.linear       | Linear                   | 3.1 K \n",
+      "30 | model.gnn.out_linear.batch_norm   | BatchNorm1d              | 64    \n",
+      "31 | model.gnn.out_linear.dropout      | Dropout                  | 0     \n",
+      "32 | model.post_nn                     | FeedForwardNN            | 1.2 K \n",
+      "33 | model.post_nn.activation          | ReLU                     | 0     \n",
+      "34 | model.post_nn.residual_layer      | ResidualConnectionNone   | 0     \n",
+      "35 | model.post_nn.layers              | ModuleList               | 1.2 K \n",
+      "36 | model.post_nn.layers.0            | FCLayer                  | 1.1 K \n",
+      "37 | model.post_nn.layers.1            | FCLayer                  | 33    \n",
+      "38 | loss_fun                          | MSELoss                  | 0     \n",
+      "--------------------------------------------------------------------------------\n",
+      "74.4 K    Trainable params\n",
       "0         Non-trainable params\n",
-      "69.7 K    Total params\n",
-      "0.279     Total estimated model params size (MB)\n",
+      "74.4 K    Total params\n",
+      "0.298     Total estimated model params size (MB)\n",
+      "2021-06-21 12:40:03.929 | INFO     | goli.data.datamodule:_load_from_cache:605 - Datamodule correctly reloaded from cache.\n",
+      "C:\\Users\\Domin\\miniconda3\\envs\\goli\\lib\\site-packages\\pytorch_lightning\\utilities\\distributed.py:68: RuntimeWarning: Found unsupported keys in the lr scheduler dict: ['mode']\n",
+      "  warnings.warn(*args, **kwargs)\n",
       "\n",
       "  | Name     | Type           | Params\n",
       "--------------------------------------------\n",
-      "0 | model    | FullDGLNetwork | 69.7 K\n",
+      "0 | model    | FullDGLNetwork | 74.4 K\n",
       "1 | loss_fun | MSELoss        | 0     \n",
       "--------------------------------------------\n",
-      "69.7 K    Trainable params\n",
+      "74.4 K    Trainable params\n",
       "0         Non-trainable params\n",
-      "69.7 K    Total params\n",
-      "0.279     Total estimated model params size (MB)\n"
+      "74.4 K    Total params\n",
+      "0.298     Total estimated model params size (MB)\n",
+      "Epoch 0:  71%|███████▏  | 5/7 [00:02<00:00,  2.29it/s, loss=3.97, v_num=15, loss/val=3.110, mae/val=1.390, pearsonr/val=-2.68e-6, f1 > 3/val=0.000, precision > 3/val=0.995]\n",
+      "Validating: 0it [00:00, ?it/s]\u001b[A\n",
+      "Validating:   0%|          | 0/2 [00:00<?, ?it/s]\u001b[A\n",
+      "Epoch 0: 100%|██████████| 7/7 [00:02<00:00,  2.87it/s, loss=3.97, v_num=15, loss/val=3.110, mae/val=1.390, pearsonr/val=-2.68e-6, f1 > 3/val=0.000, precision > 3/val=0.995]\n",
+      "Epoch 0: 100%|██████████| 7/7 [00:02<00:00,  2.67it/s, loss=3.97, v_num=15, loss/val=5.480, mae/val=1.920, pearsonr/val=-.244, f1 > 3/val=0.000, precision > 3/val=0.995]   \n",
+      "Epoch 1:  86%|████████▌ | 6/7 [00:01<00:00,  3.02it/s, loss=3.62, v_num=15, loss/val=5.480, mae/val=1.920, pearsonr/val=-.244, f1 > 3/val=0.000, precision > 3/val=0.995]\n",
+      "Validating: 0it [00:00, ?it/s]\u001b[A\n",
+      "Validating:   0%|          | 0/2 [00:00<?, ?it/s]\u001b[A\n",
+      "Validating:  50%|█████     | 1/2 [00:00<00:00,  4.35it/s]\u001b[A\n",
+      "Epoch 1: 100%|██████████| 7/7 [00:02<00:00,  2.92it/s, loss=3.62, v_num=15, loss/val=2.980, mae/val=1.370, pearsonr/val=0.137, f1 > 3/val=0.000, precision > 3/val=0.995]\n",
+      "Epoch 2:  86%|████████▌ | 6/7 [00:02<00:00,  2.61it/s, loss=3.32, v_num=15, loss/val=2.980, mae/val=1.370, pearsonr/val=0.137, f1 > 3/val=0.000, precision > 3/val=0.995]\n",
+      "Validating: 0it [00:00, ?it/s]\u001b[A\n",
+      "Validating:   0%|          | 0/2 [00:00<?, ?it/s]\u001b[A\n",
+      "Validating:  50%|█████     | 1/2 [00:00<00:00,  4.30it/s]\u001b[A\n",
+      "Epoch 2: 100%|██████████| 7/7 [00:02<00:00,  2.57it/s, loss=3.32, v_num=15, loss/val=3.850, mae/val=1.490, pearsonr/val=0.457, f1 > 3/val=0.000, precision > 3/val=0.995]\n",
+      "Epoch 3:  86%|████████▌ | 6/7 [00:01<00:00,  3.04it/s, loss=3.09, v_num=15, loss/val=3.850, mae/val=1.490, pearsonr/val=0.457, f1 > 3/val=0.000, precision > 3/val=0.995]\n",
+      "Validating: 0it [00:00, ?it/s]\u001b[A\n",
+      "Validating:   0%|          | 0/2 [00:00<?, ?it/s]\u001b[A\n",
+      "Validating:  50%|█████     | 1/2 [00:00<00:00,  4.27it/s]\u001b[A\n",
+      "Epoch 3: 100%|██████████| 7/7 [00:02<00:00,  2.92it/s, loss=3.09, v_num=15, loss/val=7.770, mae/val=2.380, pearsonr/val=0.548, f1 > 3/val=0.0476, precision > 3/val=0.800]\n",
+      "Epoch 4:  86%|████████▌ | 6/7 [00:02<00:00,  2.84it/s, loss=2.56, v_num=15, loss/val=7.770, mae/val=2.380, pearsonr/val=0.548, f1 > 3/val=0.0476, precision > 3/val=0.800]\n",
+      "Validating: 0it [00:00, ?it/s]\u001b[A\n",
+      "Validating:   0%|          | 0/2 [00:00<?, ?it/s]\u001b[A\n",
+      "Validating:  50%|█████     | 1/2 [00:00<00:00,  4.20it/s]\u001b[A\n",
+      "Epoch 4: 100%|██████████| 7/7 [00:02<00:00,  2.74it/s, loss=2.56, v_num=15, loss/val=2.720, mae/val=1.250, pearsonr/val=0.693, f1 > 3/val=0.167, precision > 3/val=0.950] \n",
+      "Epoch 5:  86%|████████▌ | 6/7 [00:02<00:00,  2.72it/s, loss=2.14, v_num=15, loss/val=2.720, mae/val=1.250, pearsonr/val=0.693, f1 > 3/val=0.167, precision > 3/val=0.950]\n",
+      "Validating: 0it [00:00, ?it/s]\u001b[A\n",
+      "Validating:   0%|          | 0/2 [00:00<?, ?it/s]\u001b[A\n",
+      "Validating:  50%|█████     | 1/2 [00:00<00:00,  2.76it/s]\u001b[A\n",
+      "Epoch 5: 100%|██████████| 7/7 [00:02<00:00,  2.56it/s, loss=2.14, v_num=15, loss/val=2.130, mae/val=1.140, pearsonr/val=0.753, f1 > 3/val=0.000, precision > 3/val=0.990]\n",
+      "Epoch 6:  86%|████████▌ | 6/7 [00:02<00:00,  2.86it/s, loss=1.78, v_num=15, loss/val=2.130, mae/val=1.140, pearsonr/val=0.753, f1 > 3/val=0.000, precision > 3/val=0.990]\n",
+      "Validating: 0it [00:00, ?it/s]\u001b[A\n",
+      "Validating:   0%|          | 0/2 [00:00<?, ?it/s]\u001b[A\n",
+      "Validating:  50%|█████     | 1/2 [00:00<00:00,  3.94it/s]\u001b[A\n",
+      "Epoch 6: 100%|██████████| 7/7 [00:02<00:00,  2.72it/s, loss=1.78, v_num=15, loss/val=1.780, mae/val=1.120, pearsonr/val=0.828, f1 > 3/val=0.000, precision > 3/val=0.995]\n",
+      "Epoch 7:  86%|████████▌ | 6/7 [00:02<00:00,  2.83it/s, loss=1.47, v_num=15, loss/val=1.780, mae/val=1.120, pearsonr/val=0.828, f1 > 3/val=0.000, precision > 3/val=0.995]\n",
+      "Validating: 0it [00:00, ?it/s]\u001b[A\n",
+      "Validating:   0%|          | 0/2 [00:00<?, ?it/s]\u001b[A\n",
+      "Validating:  50%|█████     | 1/2 [00:00<00:00,  3.97it/s]\u001b[A\n",
+      "Epoch 7: 100%|██████████| 7/7 [00:02<00:00,  2.73it/s, loss=1.47, v_num=15, loss/val=2.880, mae/val=1.510, pearsonr/val=0.861, f1 > 3/val=0.000, precision > 3/val=0.995]\n",
+      "Epoch 8:  86%|████████▌ | 6/7 [00:02<00:00,  2.82it/s, loss=1.23, v_num=15, loss/val=2.880, mae/val=1.510, pearsonr/val=0.861, f1 > 3/val=0.000, precision > 3/val=0.995]\n",
+      "Validating: 0it [00:00, ?it/s]\u001b[A\n",
+      "Validating:   0%|          | 0/2 [00:00<?, ?it/s]\u001b[A\n",
+      "Validating:  50%|█████     | 1/2 [00:00<00:00,  4.41it/s]\u001b[A\n",
+      "Epoch 8: 100%|██████████| 7/7 [00:02<00:00,  2.73it/s, loss=1.23, v_num=15, loss/val=6.910, mae/val=2.470, pearsonr/val=0.842, f1 > 3/val=0.000, precision > 3/val=0.995]\n",
+      "Epoch 9:  86%|████████▌ | 6/7 [00:02<00:00,  2.89it/s, loss=1.05, v_num=15, loss/val=6.910, mae/val=2.470, pearsonr/val=0.842, f1 > 3/val=0.000, precision > 3/val=0.995]\n",
+      "Validating: 0it [00:00, ?it/s]\u001b[A\n",
+      "Validating:   0%|          | 0/2 [00:00<?, ?it/s]\u001b[A\n",
+      "Validating:  50%|█████     | 1/2 [00:00<00:00,  3.98it/s]\u001b[A\n",
+      "Epoch 9: 100%|██████████| 7/7 [00:02<00:00,  2.76it/s, loss=1.05, v_num=15, loss/val=5.230, mae/val=2.150, pearsonr/val=0.890, f1 > 3/val=0.000, precision > 3/val=0.995]\n",
+      "Epoch 10:  86%|████████▌ | 6/7 [00:02<00:00,  2.81it/s, loss=0.918, v_num=15, loss/val=5.230, mae/val=2.150, pearsonr/val=0.890, f1 > 3/val=0.000, precision > 3/val=0.995]\n",
+      "Validating: 0it [00:00, ?it/s]\u001b[A\n",
+      "Validating:   0%|          | 0/2 [00:00<?, ?it/s]\u001b[A\n",
+      "Validating:  50%|█████     | 1/2 [00:00<00:00,  4.12it/s]\u001b[A\n",
+      "Epoch 10: 100%|██████████| 7/7 [00:02<00:00,  2.71it/s, loss=0.918, v_num=15, loss/val=9.860, mae/val=3.030, pearsonr/val=0.869, f1 > 3/val=0.000, precision > 3/val=0.995]\n",
+      "Epoch 11:  86%|████████▌ | 6/7 [00:02<00:00,  2.81it/s, loss=0.805, v_num=15, loss/val=9.860, mae/val=3.030, pearsonr/val=0.869, f1 > 3/val=0.000, precision > 3/val=0.995]\n",
+      "Validating: 0it [00:00, ?it/s]\u001b[A\n",
+      "Validating:   0%|          | 0/2 [00:00<?, ?it/s]\u001b[A\n",
+      "Validating:  50%|█████     | 1/2 [00:00<00:00,  3.91it/s]\u001b[A\n",
+      "Epoch 11: 100%|██████████| 7/7 [00:02<00:00,  2.70it/s, loss=0.805, v_num=15, loss/val=4.080, mae/val=1.900, pearsonr/val=0.899, f1 > 3/val=0.000, precision > 3/val=0.995]\n",
+      "Epoch 12:  86%|████████▌ | 6/7 [00:02<00:00,  2.76it/s, loss=0.736, v_num=15, loss/val=4.080, mae/val=1.900, pearsonr/val=0.899, f1 > 3/val=0.000, precision > 3/val=0.995]\n",
+      "Validating: 0it [00:00, ?it/s]\u001b[A\n",
+      "Validating:   0%|          | 0/2 [00:00<?, ?it/s]\u001b[A\n",
+      "Validating:  50%|█████     | 1/2 [00:00<00:00,  3.46it/s]\u001b[A\n",
+      "Epoch 12: 100%|██████████| 7/7 [00:02<00:00,  2.62it/s, loss=0.736, v_num=15, loss/val=1.540, mae/val=1.100, pearsonr/val=0.934, f1 > 3/val=0.000, precision > 3/val=0.995]\n",
+      "Epoch 13:  86%|████████▌ | 6/7 [00:02<00:00,  2.77it/s, loss=0.694, v_num=15, loss/val=1.540, mae/val=1.100, pearsonr/val=0.934, f1 > 3/val=0.000, precision > 3/val=0.995]\n",
+      "Validating: 0it [00:00, ?it/s]\u001b[A\n",
+      "Validating:   0%|          | 0/2 [00:00<?, ?it/s]\u001b[A\n",
+      "Validating:  50%|█████     | 1/2 [00:00<00:00,  3.71it/s]\u001b[A\n",
+      "Epoch 13: 100%|██████████| 7/7 [00:02<00:00,  2.63it/s, loss=0.694, v_num=15, loss/val=1.140, mae/val=0.928, pearsonr/val=0.927, f1 > 3/val=0.000, precision > 3/val=0.995]\n",
+      "Epoch 14:  86%|████████▌ | 6/7 [00:02<00:00,  2.75it/s, loss=0.662, v_num=15, loss/val=1.140, mae/val=0.928, pearsonr/val=0.927, f1 > 3/val=0.000, precision > 3/val=0.995]\n",
+      "Validating: 0it [00:00, ?it/s]\u001b[A\n",
+      "Validating:   0%|          | 0/2 [00:00<?, ?it/s]\u001b[A\n",
+      "Validating:  50%|█████     | 1/2 [00:00<00:00,  4.29it/s]\u001b[A\n",
+      "Epoch 14: 100%|██████████| 7/7 [00:02<00:00,  2.67it/s, loss=0.662, v_num=15, loss/val=0.830, mae/val=0.711, pearsonr/val=0.908, f1 > 3/val=0.000, precision > 3/val=0.995]\n",
+      "Epoch 15:  86%|████████▌ | 6/7 [00:02<00:00,  2.80it/s, loss=0.658, v_num=15, loss/val=0.830, mae/val=0.711, pearsonr/val=0.908, f1 > 3/val=0.000, precision > 3/val=0.995]\n",
+      "Validating: 0it [00:00, ?it/s]\u001b[A\n",
+      "Validating:   0%|          | 0/2 [00:00<?, ?it/s]\u001b[A\n",
+      "Validating:  50%|█████     | 1/2 [00:00<00:00,  3.95it/s]\u001b[A\n",
+      "Epoch 15: 100%|██████████| 7/7 [00:02<00:00,  2.69it/s, loss=0.658, v_num=15, loss/val=0.476, mae/val=0.527, pearsonr/val=0.930, f1 > 3/val=0.000, precision > 3/val=0.995]\n",
+      "Epoch 16:  86%|████████▌ | 6/7 [00:02<00:00,  2.83it/s, loss=0.64, v_num=15, loss/val=0.476, mae/val=0.527, pearsonr/val=0.930, f1 > 3/val=0.000, precision > 3/val=0.995]\n",
+      "Validating: 0it [00:00, ?it/s]\u001b[A\n",
+      "Validating:   0%|          | 0/2 [00:00<?, ?it/s]\u001b[A\n",
+      "Validating:  50%|█████     | 1/2 [00:00<00:00,  4.02it/s]\u001b[A\n",
+      "Epoch 16: 100%|██████████| 7/7 [00:02<00:00,  2.72it/s, loss=0.64, v_num=15, loss/val=0.780, mae/val=0.727, pearsonr/val=0.924, f1 > 3/val=0.400, precision > 3/val=0.985]\n",
+      "Epoch 17:  86%|████████▌ | 6/7 [00:02<00:00,  2.83it/s, loss=0.594, v_num=15, loss/val=0.780, mae/val=0.727, pearsonr/val=0.924, f1 > 3/val=0.400, precision > 3/val=0.985]\n",
+      "Validating: 0it [00:00, ?it/s]\u001b[A\n",
+      "Validating:   0%|          | 0/2 [00:00<?, ?it/s]\u001b[A\n",
+      "Validating:  50%|█████     | 1/2 [00:00<00:00,  3.98it/s]\u001b[A\n",
+      "Epoch 17: 100%|██████████| 7/7 [00:02<00:00,  2.72it/s, loss=0.594, v_num=15, loss/val=0.593, mae/val=0.624, pearsonr/val=0.934, f1 > 3/val=0.000, precision > 3/val=0.995]\n",
+      "Epoch 18:  86%|████████▌ | 6/7 [00:02<00:00,  2.86it/s, loss=0.559, v_num=15, loss/val=0.593, mae/val=0.624, pearsonr/val=0.934, f1 > 3/val=0.000, precision > 3/val=0.995]\n",
+      "Validating: 0it [00:00, ?it/s]\u001b[A\n",
+      "Validating:   0%|          | 0/2 [00:00<?, ?it/s]\u001b[A\n",
+      "Validating:  50%|█████     | 1/2 [00:00<00:00,  3.77it/s]\u001b[A\n",
+      "Epoch 18: 100%|██████████| 7/7 [00:02<00:00,  2.70it/s, loss=0.559, v_num=15, loss/val=0.846, mae/val=0.773, pearsonr/val=0.930, f1 > 3/val=0.333, precision > 3/val=0.980]\n",
+      "Epoch 19:  86%|████████▌ | 6/7 [00:02<00:00,  2.84it/s, loss=0.515, v_num=15, loss/val=0.846, mae/val=0.773, pearsonr/val=0.930, f1 > 3/val=0.333, precision > 3/val=0.980]\n",
+      "Validating: 0it [00:00, ?it/s]\u001b[A\n",
+      "Validating:   0%|          | 0/2 [00:00<?, ?it/s]\u001b[A\n",
+      "Validating:  50%|█████     | 1/2 [00:00<00:00,  3.85it/s]\u001b[A\n",
+      "Epoch 19: 100%|██████████| 7/7 [00:02<00:00,  2.72it/s, loss=0.515, v_num=15, loss/val=0.513, mae/val=0.572, pearsonr/val=0.928, f1 > 3/val=0.000, precision > 3/val=0.990]\n",
+      "Epoch 20:  86%|████████▌ | 6/7 [00:02<00:00,  2.81it/s, loss=0.491, v_num=15, loss/val=0.513, mae/val=0.572, pearsonr/val=0.928, f1 > 3/val=0.000, precision > 3/val=0.990]\n",
+      "Validating: 0it [00:00, ?it/s]\u001b[A\n",
+      "Validating:   0%|          | 0/2 [00:00<?, ?it/s]\u001b[A\n",
+      "Validating:  50%|█████     | 1/2 [00:00<00:00,  3.61it/s]\u001b[A\n",
+      "Epoch 20: 100%|██████████| 7/7 [00:02<00:00,  2.67it/s, loss=0.491, v_num=15, loss/val=0.559, mae/val=0.580, pearsonr/val=0.931, f1 > 3/val=0.000, precision > 3/val=0.995]\n",
+      "Epoch 21:  86%|████████▌ | 6/7 [00:02<00:00,  2.81it/s, loss=0.485, v_num=15, loss/val=0.559, mae/val=0.580, pearsonr/val=0.931, f1 > 3/val=0.000, precision > 3/val=0.995]\n",
+      "Validating: 0it [00:00, ?it/s]\u001b[A\n",
+      "Validating:   0%|          | 0/2 [00:00<?, ?it/s]\u001b[A\n",
+      "Validating:  50%|█████     | 1/2 [00:00<00:00,  4.15it/s]\u001b[A\n",
+      "Epoch 21: 100%|██████████| 7/7 [00:02<00:00,  2.72it/s, loss=0.485, v_num=15, loss/val=0.380, mae/val=0.473, pearsonr/val=0.939, f1 > 3/val=0.667, precision > 3/val=0.995]\n",
+      "Epoch 22:  86%|████████▌ | 6/7 [00:02<00:00,  2.80it/s, loss=0.471, v_num=15, loss/val=0.380, mae/val=0.473, pearsonr/val=0.939, f1 > 3/val=0.667, precision > 3/val=0.995]\n",
+      "Validating: 0it [00:00, ?it/s]\u001b[A\n",
+      "Validating:   0%|          | 0/2 [00:00<?, ?it/s]\u001b[A\n",
+      "Validating:  50%|█████     | 1/2 [00:00<00:00,  4.04it/s]\u001b[A\n",
+      "Epoch 22: 100%|██████████| 7/7 [00:02<00:00,  2.68it/s, loss=0.471, v_num=15, loss/val=0.340, mae/val=0.433, pearsonr/val=0.942, f1 > 3/val=0.000, precision > 3/val=0.995]\n",
+      "Epoch 23:  86%|████████▌ | 6/7 [00:02<00:00,  2.84it/s, loss=0.467, v_num=15, loss/val=0.340, mae/val=0.433, pearsonr/val=0.942, f1 > 3/val=0.000, precision > 3/val=0.995]\n",
+      "Validating: 0it [00:00, ?it/s]\u001b[A\n",
+      "Validating:   0%|          | 0/2 [00:00<?, ?it/s]\u001b[A\n",
+      "Validating:  50%|█████     | 1/2 [00:00<00:00,  3.91it/s]\u001b[A\n",
+      "Epoch 23: 100%|██████████| 7/7 [00:02<00:00,  2.72it/s, loss=0.467, v_num=15, loss/val=0.423, mae/val=0.501, pearsonr/val=0.938, f1 > 3/val=0.667, precision > 3/val=0.995]\n",
+      "Epoch 24:  86%|████████▌ | 6/7 [00:02<00:00,  2.82it/s, loss=0.473, v_num=15, loss/val=0.423, mae/val=0.501, pearsonr/val=0.938, f1 > 3/val=0.667, precision > 3/val=0.995]\n",
+      "Validating: 0it [00:00, ?it/s]\u001b[A\n",
+      "Validating:   0%|          | 0/2 [00:00<?, ?it/s]\u001b[A\n",
+      "Validating:  50%|█████     | 1/2 [00:00<00:00,  4.18it/s]\u001b[A\n",
+      "Epoch 24: 100%|██████████| 7/7 [00:02<00:00,  2.71it/s, loss=0.473, v_num=15, loss/val=0.326, mae/val=0.430, pearsonr/val=0.946, f1 > 3/val=0.000, precision > 3/val=0.995]\n",
+      "Epoch 24: 100%|██████████| 7/7 [00:02<00:00,  2.65it/s, loss=0.473, v_num=15, loss/val=0.326, mae/val=0.430, pearsonr/val=0.946, f1 > 3/val=0.000, precision > 3/val=0.995]\n"
      ]
     },
     {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "61dc1894ee264599ab493d982b390430",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation sanity check: 0it [00:00, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "/home/dominique/anaconda3/envs/goli/lib/python3.8/site-packages/pytorch_lightning/utilities/distributed.py:51: UserWarning: The validation_epoch_end should not return anything as of 9.1. To log, use self.log(...) or self.write(...) directly in the LightningModule\n",
-      "  warnings.warn(*args, **kwargs)\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "657567b3d0b546a1a648173c2bfb1e4a",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Training: 0it [00:00, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "01792491c7fd49b08ce5086832135c7b",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validating: 0it [00:00, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "b9d96f2469234fe380d07ffd806350ed",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validating: 0it [00:00, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "eed00b1f81524fe99e07e2084c952532",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validating: 0it [00:00, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "c435a81142c24be09a0e38d7575b365b",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validating: 0it [00:00, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "ae0357d542ec4021b0c9c38fb38bd11c",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validating: 0it [00:00, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "a1d24f83f1504d3b80b0741d1c0f404b",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validating: 0it [00:00, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "647aee1f810f407c90697c394d0f604f",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validating: 0it [00:00, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "c69086a4802e421f816cab1ebbc20ae9",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validating: 0it [00:00, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "ebb4a32a0f78470ba21ff5bea8b450f4",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validating: 0it [00:00, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "840d27d095344fcd9a74862e61a2fe7f",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validating: 0it [00:00, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "905c7ee70f4b4282871c130b0c7b9f0a",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validating: 0it [00:00, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "daab3285c0854c23a4e3dd47846c820e",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validating: 0it [00:00, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "e0e0d9096fc64198a470ae1b3cd7f351",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validating: 0it [00:00, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "99ac351f4e334e8c838a6913ef6bee08",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validating: 0it [00:00, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "69b47fad071248eab8095d67e33b5d5e",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validating: 0it [00:00, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "8c68dcc01135429e845427bb6908f414",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validating: 0it [00:00, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "096aaea9ce2649fba9bf70b99b7e7955",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validating: 0it [00:00, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "d8ecff999d934a119157a3e0ca7a1c6a",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validating: 0it [00:00, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "4287a56d059b4eb2966eb2e90498a210",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validating: 0it [00:00, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "0a5ffab4db4e4768a4876b01a8b10f96",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validating: 0it [00:00, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "6177ef595f9542598e5b065d6d77bb32",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validating: 0it [00:00, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "e86938e35b0b443791119e37dd2e2199",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validating: 0it [00:00, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "98aff21b49cc434dbaaaf12c355ab783",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validating: 0it [00:00, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "e4c88c49c0c843c09934e9786e9b6aa5",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validating: 0it [00:00, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "dcb917418a084d4ba36d57f5b0406819",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validating: 0it [00:00, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
+     "output_type": "execute_result",
      "data": {
       "text/plain": [
        "1"
       ]
      },
-     "execution_count": 10,
      "metadata": {},
-     "output_type": "execute_result"
+     "execution_count": 10
     }
    ],
    "source": [
@@ -901,33 +533,87 @@
     "    in_dim_edges=datamodule.num_edge_feats,\n",
     ")\n",
     "\n",
+    "# Load and print the metrics\n",
     "metrics = load_metrics(cfg)\n",
     "print(metrics)\n",
     "\n",
+    "# Load the predictor, print the model, and print a summary of the number of parameters\n",
     "predictor = load_predictor(cfg, model_class, model_kwargs, metrics)\n",
-    "\n",
     "print(predictor.model)\n",
     "print(predictor.summarize(mode=4, to_print=False))\n",
     "\n",
-    "trainer = load_trainer(cfg, metrics)\n",
-    "\n",
-    "# Run the model training\n",
+    "# Load the trainer, and start the training\n",
+    "trainer = load_trainer(cfg)\n",
     "trainer.fit(model=predictor, datamodule=datamodule)"
    ]
   },
+  {
+   "source": [
+    "## Testing the model\n",
+    "Once the model is trained, we can use the same datamodule to get the results on the test set. Here, `ckpt_path` refers to the checkpoint path where the model at the best validation step was saved. Thus, the results on the test set represent the early stopping.\n",
+    "\n",
+    "All the metrics that were computed on the validation set are then computed on the test set, printed, and saved into the `metrics.yaml` file."
+   ],
+   "cell_type": "markdown",
+   "metadata": {}
+  },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 11,
    "metadata": {},
-   "outputs": [],
-   "source": []
+   "outputs": [
+    {
+     "output_type": "stream",
+     "name": "stdout",
+     "text": [
+      "Testing: 100%|██████████| 2/2 [00:00<00:00,  3.23it/s]\n",
+      "--------------------------------------------------------------------------------\n",
+      "DATALOADER:0 TEST RESULTS\n",
+      "{'MSELoss/test': 0.6393043398857117,\n",
+      " 'f1 > 3/test': 0.0,\n",
+      " 'f1 > 5/test': 0.0,\n",
+      " 'loss/test': 0.6393043398857117,\n",
+      " 'mae/test': 0.5762802362442017,\n",
+      " 'mean_pred/test': -0.48094016313552856,\n",
+      " 'mean_target/test': -0.6447566151618958,\n",
+      " 'pearsonr/test': 0.9359009265899658,\n",
+      " 'precision > 3/test': 0.9900000095367432,\n",
+      " 'std_pred/test': 1.7718788385391235,\n",
+      " 'std_target/test': 2.1336562633514404}\n",
+      "--------------------------------------------------------------------------------\n"
+     ]
+    },
+    {
+     "output_type": "execute_result",
+     "data": {
+      "text/plain": [
+       "[{'mean_pred/test': -0.48094016313552856,\n",
+       "  'std_pred/test': 1.7718788385391235,\n",
+       "  'mean_target/test': -0.6447566151618958,\n",
+       "  'std_target/test': 2.1336562633514404,\n",
+       "  'mae/test': 0.5762802362442017,\n",
+       "  'pearsonr/test': 0.9359009265899658,\n",
+       "  'f1 > 3/test': 0.0,\n",
+       "  'f1 > 5/test': 0.0,\n",
+       "  'precision > 3/test': 0.9900000095367432,\n",
+       "  'MSELoss/test': 0.6393043398857117,\n",
+       "  'loss/test': 0.6393043398857117}]"
+      ]
+     },
+     "metadata": {},
+     "execution_count": 11
+    }
+   ],
+   "source": [
+    "ckpt_path = trainer.checkpoint_callbacks[0].best_model_path\n",
+    "trainer.test(model=predictor, datamodule=datamodule, ckpt_path=ckpt_path)"
+   ]
   }
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3",
-   "language": "python",
-   "name": "python3"
+   "name": "python3",
+   "display_name": "Python 3.8.8 64-bit ('goli': conda)"
   },
   "language_info": {
    "codemirror_mode": {
@@ -939,7 +625,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.4"
+   "version": "3.8.8"
   },
   "widgets": {
    "application/vnd.jupyter.widget-state+json": {
@@ -947,8 +633,11 @@
     "version_major": 2,
     "version_minor": 0
    }
+  },
+  "interpreter": {
+   "hash": "f4a99d018a205fcbcc0480c84566beaebcb91b08d0414b39a842df533e2a1d25"
   }
  },
  "nbformat": 4,
  "nbformat_minor": 4
-}
+}
\ No newline at end of file
diff --git a/env.yml b/env.yml
index 7aef9a329..1d5b257f8 100644
--- a/env.yml
+++ b/env.yml
@@ -4,7 +4,7 @@ channels:
   - conda-forge
 
 dependencies:
-  - python >=3.7
+  - python >=3.7, <3.9
   - pip
   - click
   - loguru
@@ -23,9 +23,9 @@ dependencies:
   - seaborn
 
   # cloud IO
-  - fsspec
-  - s3fs
-  - gcsfs
+  - fsspec >=2021.6
+  - s3fs >=2021.6
+  - gcsfs >=2021.6
   - appdirs
 
   # ML packages
diff --git a/expts/config_ZINC_bench_gnn.yaml b/expts/config_ZINC_bench_gnn.yaml
index 76db1d973..a5782ffcc 100644
--- a/expts/config_ZINC_bench_gnn.yaml
+++ b/expts/config_ZINC_bench_gnn.yaml
@@ -4,7 +4,7 @@ constants:
 
 datamodule:
   module_type: "DGLFromSmilesDataModule"
-    df_path: goli/data/ZINC_bench_gnn/smiles_score.csv
+    df_path: https://storage.googleapis.com/goli-public/datasets/goli-zinc-bench-gnn.zip
     cache_data_path: goli/data/cache/ZINC_bench_gnn/smiles_score.cache
     label_cols: ['score']
     smiles_col: SMILES
@@ -60,7 +60,7 @@ architecture:
     activation: relu
     last_activation: none
     dropout: *dropout
-    batch_norm: *batch_norm 
+    batch_norm: *batch_norm
     last_batch_norm: *batch_norm
     residual_type: simple
     pooling: 'sum'
@@ -70,7 +70,7 @@ architecture:
       # num_heads: 3
       aggregators: [mean, max, dir1/dx_abs, dir1/smooth]
       scalers: [identity, amplification, attenuation]
-  
+
   post_nn:
     out_dim: 1
     hidden_dims: *middle_dim
@@ -137,7 +137,7 @@ metrics:
     threshold_kwargs: *threshold_0
 
 trainer:
-  logger: 
+  logger:
     save_dir: logs/ZINC_bench_gnn
   early_stopping:
     monitor: *monitor
@@ -150,10 +150,9 @@ trainer:
     monitor: *monitor
     mode: *mode
     save_top_k: 1
-    period: 1
+    every_n_val_epochs: 1
   trainer:
     max_epochs: 2000
     min_epochs: 100
     gpus: 1
 
-    
\ No newline at end of file
diff --git a/expts/config_micro_ZINC.yaml b/expts/config_micro_ZINC.yaml
index f67018daa..5bf79823a 100644
--- a/expts/config_micro_ZINC.yaml
+++ b/expts/config_micro_ZINC.yaml
@@ -6,7 +6,7 @@ datamodule:
   module_type: "DGLFromSmilesDataModule"
   args:
     df_path: goli/data/micro_ZINC/micro_ZINC.csv
-    cache_data_path: goli/data/cache/micro_ZINC/full.cache
+    cache_data_path: goli/data/cache/micro_ZINC/full.cache.gz
     label_cols: ['score']
     smiles_col: SMILES
 
@@ -72,7 +72,7 @@ architecture:
     activation: relu
     last_activation: none
     dropout: *dropout
-    batch_norm: *batch_norm 
+    batch_norm: *batch_norm
     last_batch_norm: *batch_norm
     residual_type: simple
     pooling: [sum, max, dir1]
@@ -82,7 +82,7 @@ architecture:
       # num_heads: 3
       aggregators: [mean, max, dir1/dx_abs, dir1/smooth]
       scalers: [identity, amplification, attenuation]
-  
+
   post_nn:
     out_dim: 1
     hidden_dims: 32
@@ -160,10 +160,9 @@ trainer:
     monitor: *monitor
     mode: *mode
     save_top_k: 1
-    period: 1
+    every_n_val_epochs: 1
   trainer:
     max_epochs: 25
     min_epochs: 5
     gpus: 1
 
-    
\ No newline at end of file
diff --git a/expts/config_molHIV.yaml b/expts/config_molHIV.yaml
index bd5dd2095..f920a7508 100644
--- a/expts/config_molHIV.yaml
+++ b/expts/config_molHIV.yaml
@@ -7,7 +7,7 @@ datamodule:
   args:
     cache_data_path: null
     dataset_name: "ogbg-molhiv"
-    cache_data_path: goli/data/cache/ogb-molhiv/full.cache
+    cache_data_path: goli/data/cache/ogb-molhiv/full.cache.gz
 
     # Weights
     weights_type: sample_balanced
@@ -59,7 +59,7 @@ architecture:
     activation: none
     last_activation: none
     dropout: *dropout
-    batch_norm: *batch_norm 
+    batch_norm: *batch_norm
     last_batch_norm: *batch_norm
     residual_type: simple
     pooling: sum
@@ -69,7 +69,7 @@ architecture:
       # num_heads: 3
       aggregators: [mean, max, min, dir1/dx_abs, dir1/smooth]
       scalers: [identity]
-  
+
   post_nn:
     out_dim: 1
     hidden_dims: *hidden_dim
@@ -156,10 +156,9 @@ trainer:
     monitor: *monitor
     mode: *mode
     save_top_k: 1
-    period: 1
+    every_n_val_epochs: 1
   trainer:
     max_epochs: 1000
     min_epochs: 100
     gpus: 1
 
-    
\ No newline at end of file
diff --git a/expts/config_molHIV_pretrained.yaml b/expts/config_molHIV_pretrained.yaml
index 12170c60b..09ee7694f 100644
--- a/expts/config_molHIV_pretrained.yaml
+++ b/expts/config_molHIV_pretrained.yaml
@@ -7,7 +7,7 @@ datamodule:
   args:
     cache_data_path: null
     dataset_name: "ogbg-molhiv"
-    cache_data_path: goli/data/cache/ogb-molhiv/full.cache
+    cache_data_path: goli/data/cache/ogb-molhiv/full.cache.gz
 
     # Weights
     weights_type: sample_balanced
diff --git a/expts/config_molPCBA.yaml b/expts/config_molPCBA.yaml
index 548d183b4..02966b37b 100644
--- a/expts/config_molPCBA.yaml
+++ b/expts/config_molPCBA.yaml
@@ -73,7 +73,7 @@ architecture:
     last_activation: none
     dropout: &dropout_gnn 0.2
     last_dropout: *dropout_gnn
-    batch_norm: *batch_norm 
+    batch_norm: *batch_norm
     last_batch_norm: *batch_norm
     residual_type: simple
     pooling: ['sum', 'max']
@@ -83,7 +83,7 @@ architecture:
       # num_heads: 3
       aggregators: [mean, max, sum, dir1/dx_abs]
       scalers: [identity]
-  
+
   post_nn:
     out_dim: 128
     hidden_dims: *hidden_dim
@@ -187,11 +187,10 @@ trainer:
     monitor: *monitor
     mode: *mode
     save_top_k: 1
-    period: 1
+    every_n_val_epochs: 1
   trainer:
     max_epochs: 1000
     min_epochs: 100
     gpus: 1
     accumulate_grad_batches: 1
 
-    
\ No newline at end of file
diff --git a/expts/config_molbace_pretrained.yaml b/expts/config_molbace_pretrained.yaml
index 498bd4bba..7bba30594 100644
--- a/expts/config_molbace_pretrained.yaml
+++ b/expts/config_molbace_pretrained.yaml
@@ -7,7 +7,7 @@ datamodule:
   args:
     cache_data_path: null
     dataset_name: "ogbg-molbace"
-    cache_data_path: goli/data/cache/ogb-molbace/full.cache
+    cache_data_path: goli/data/cache/ogb-molbace/full.cache.gz
 
     # Weights
     weights_type: null
diff --git a/expts/config_mollipo_pretrained.yaml b/expts/config_mollipo_pretrained.yaml
index 20941d092..1a9663ed3 100644
--- a/expts/config_mollipo_pretrained.yaml
+++ b/expts/config_mollipo_pretrained.yaml
@@ -7,7 +7,7 @@ datamodule:
   args:
     cache_data_path: null
     dataset_name: "ogbg-mollipo"
-    cache_data_path: goli/data/cache/ogb-mollipo/full.cache
+    cache_data_path: goli/data/cache/ogb-mollipo/full.cache.gz
 
     # Weights
     weights_type: null
diff --git a/expts/config_moltox21_pretrained.yaml b/expts/config_moltox21_pretrained.yaml
index 7b8369bd3..d35ae2e71 100644
--- a/expts/config_moltox21_pretrained.yaml
+++ b/expts/config_moltox21_pretrained.yaml
@@ -7,7 +7,7 @@ datamodule:
   args:
     cache_data_path: null
     dataset_name: "ogbg-moltox21"
-    cache_data_path: goli/data/cache/ogb-moltox21/full.cache
+    cache_data_path: goli/data/cache/ogb-moltox21/full.cache.gz
 
     # Weights
     weights_type: null
diff --git a/goli/data/datamodule.py b/goli/data/datamodule.py
index 5100d968c..e0ccbfc55 100644
--- a/goli/data/datamodule.py
+++ b/goli/data/datamodule.py
@@ -547,7 +547,7 @@ def _save_to_cache(self):
             # Save featurization args used
             cache["featurization_args"] = mol_to_dglgraph_signature(dict(self.featurization or {}))
 
-            with fsspec.open(self.cache_data_path, "wb") as f:
+            with fsspec.open(self.cache_data_path, "wb", compression="infer") as f:
                 torch.save(cache, f)
 
     def _load_from_cache(self):
@@ -567,7 +567,7 @@ def _load_from_cache(self):
         logger.info(f"Try reloading the data module from {self.cache_data_path}.")
 
         # Load cache
-        with fsspec.open(self.cache_data_path, "rb") as f:
+        with fsspec.open(self.cache_data_path, "rb", compression="infer") as f:
             cache = torch.load(f)
 
         # Are the required keys present?
diff --git a/goli/trainer/metrics.py b/goli/trainer/metrics.py
index 3d505c255..dfb4f0a67 100644
--- a/goli/trainer/metrics.py
+++ b/goli/trainer/metrics.py
@@ -1,14 +1,14 @@
 from typing import Union, Callable, Optional, Dict, Any
 
 from copy import deepcopy
+
 import torch
-from torch.nn import functional as F
-import torch.nn as nn
+
 import operator as op
 
-from pytorch_lightning.metrics.utils import reduce
-from pytorch_lightning.metrics.functional import auroc
-from pytorch_lightning.metrics.functional import (
+from torchmetrics.utilities import reduce
+from torchmetrics.functional import auroc
+from torchmetrics.functional import (
     accuracy,
     average_precision,
     confusion_matrix,
@@ -18,7 +18,6 @@
     precision,
     recall,
     auroc,
-    multiclass_auroc,
     mean_absolute_error,
     mean_squared_error,
 )
@@ -181,7 +180,6 @@ def spearmanr(preds: torch.Tensor, target: torch.Tensor, reduction: str = "eleme
     "precisionrecallcurve": precision_recall_curve,
     "precision": precision,
     "recall": recall,
-    "multiclass_auroc": multiclass_auroc,
 }
 
 METRICS_REGRESSION = {
diff --git a/notebooks/dev-training-loop.ipynb b/notebooks/dev-training-loop.ipynb
index a1284b5de..0c2c1012f 100644
--- a/notebooks/dev-training-loop.ipynb
+++ b/notebooks/dev-training-loop.ipynb
@@ -69,7 +69,7 @@
      "text": [
       "datamodule:\n",
       "  df_path: goli/data/micro_ZINC/micro_ZINC.csv\n",
-      "  cache_data_path: goli/data/cache/micro_ZINC/full.cache\n",
+      "  cache_data_path: goli/data/cache/micro_ZINC/full.cache.gz\n",
       "  label_cols:\n",
       "  - score\n",
       "  smiles_col: SMILES\n",
@@ -358,7 +358,7 @@
       "  warnings.warn(*args, **kwargs)\n",
       "GPU available: False, used: False\n",
       "TPU available: None, using: 0 TPU cores\n",
-      "2021-03-25 09:44:37.331 | INFO     | goli.data.datamodule:prepare_data:153 - Reload data from goli/data/cache/micro_ZINC/full.cache.\n",
+      "2021-03-25 09:44:37.331 | INFO     | goli.data.datamodule:prepare_data:153 - Reload data from goli/data/cache/micro_ZINC/full.cache.gz.\n",
       "\n",
       "datamodule:\n",
       " name: DGLFromSmilesDataModule\n",


Test patch:
diff --git a/tests/data/config_micro_ZINC.yaml b/tests/data/config_micro_ZINC.yaml
index f67018daa..5bf79823a 100644
--- a/tests/data/config_micro_ZINC.yaml
+++ b/tests/data/config_micro_ZINC.yaml
@@ -6,7 +6,7 @@ datamodule:
   module_type: "DGLFromSmilesDataModule"
   args:
     df_path: goli/data/micro_ZINC/micro_ZINC.csv
-    cache_data_path: goli/data/cache/micro_ZINC/full.cache
+    cache_data_path: goli/data/cache/micro_ZINC/full.cache.gz
     label_cols: ['score']
     smiles_col: SMILES
 
@@ -72,7 +72,7 @@ architecture:
     activation: relu
     last_activation: none
     dropout: *dropout
-    batch_norm: *batch_norm 
+    batch_norm: *batch_norm
     last_batch_norm: *batch_norm
     residual_type: simple
     pooling: [sum, max, dir1]
@@ -82,7 +82,7 @@ architecture:
       # num_heads: 3
       aggregators: [mean, max, dir1/dx_abs, dir1/smooth]
       scalers: [identity, amplification, attenuation]
-  
+
   post_nn:
     out_dim: 1
     hidden_dims: 32
@@ -160,10 +160,9 @@ trainer:
     monitor: *monitor
     mode: *mode
     save_top_k: 1
-    period: 1
+    every_n_val_epochs: 1
   trainer:
     max_epochs: 25
     min_epochs: 5
     gpus: 1
 
-    
\ No newline at end of file