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

Issue:
compatibility with PCQM4M and PCQM4Mv2
OGB changed how they handle the LSC datasets such as PCQM4M and PCQM4Mv2, and I had to adapt the code. Some parts such as download links are hard-coded.

---



Discussion / hints:

Fix patch:
diff --git a/expts/config_micro-PCBA.yaml b/expts/config_micro-PCBA.yaml
index 49e964fdc..236b03365 100644
--- a/expts/config_micro-PCBA.yaml
+++ b/expts/config_micro-PCBA.yaml
@@ -115,8 +115,9 @@ predictor:
     mode: &mode max
     frequency: 1
   target_nan_mask: ignore-mean-label # null: no mask, 0: 0 mask, ignore-flatten: ignore nan values from loss and flatten tensor, ignore-mean-label: ignore nan values from loss and average over labels
-  flag_n_steps: 0
-
+  flag_kwargs:
+    n_steps: 0
+    alpha: 0.01
 
 metrics:
   - name: mae
diff --git a/expts/config_micro_ZINC.yaml b/expts/config_micro_ZINC.yaml
index fc29055d1..24d62c3d7 100644
--- a/expts/config_micro_ZINC.yaml
+++ b/expts/config_micro_ZINC.yaml
@@ -102,14 +102,18 @@ predictor:
   optim_kwargs:
     lr: 1.e-2
     weight_decay: 1.e-7
-  lr_reduce_on_plateau_kwargs:
+  torch_scheduler_kwargs:
+    module_type: ReduceLROnPlateau
     factor: 0.5
     patience: 7
   scheduler_kwargs:
     monitor: &monitor loss/val
+    mode: min
     frequency: 1
   target_nan_mask: 0 # null: no mask, 0: 0 mask, ignore: ignore nan values from loss
-  flag_n_steps: 0
+  flag_kwargs:
+    n_steps: 1
+    alpha: 0.01
 
 metrics:
   - name: mae
diff --git a/expts/config_molPCQM4M.yaml b/expts/config_molPCQM4M.yaml
index 5f10f0814..b167ec6ea 100644
--- a/expts/config_molPCQM4M.yaml
+++ b/expts/config_molPCQM4M.yaml
@@ -6,8 +6,7 @@ datamodule:
   module_type: "DGLOGBDataModule"
   args:
     cache_data_path: null
-    dataset_name: "ogbg-molpcqm4m"
-
+    dataset_name: "ogbg-lsc-pcqm4m"
 
     # Weights
     weights_type: null
@@ -99,7 +98,7 @@ architecture:
     residual_type: simple
 
 predictor:
-  metrics_on_progress_bar: ["mse", "mae", "pearsonr"]
+  metrics_on_progress_bar: [ "mae", "pearsonr"]
   loss_fun: mse
   random_seed: *seed
   optim_kwargs:
@@ -176,7 +175,7 @@ trainer:
     monitor: *monitor
     mode: *mode
     save_top_k: 1
-    period: 1
+    every_n_val_epochs: 1
   trainer:
     max_epochs: 1000
     min_epochs: 100
diff --git a/expts/config_molPCQM4Mv2.yaml b/expts/config_molPCQM4Mv2.yaml
new file mode 100644
index 000000000..96dd10145
--- /dev/null
+++ b/expts/config_molPCQM4Mv2.yaml
@@ -0,0 +1,184 @@
+constants:
+  seed: &seed 42
+  raise_train_error: True   # Whether the code should raise an error if it crashes during training
+
+datamodule:
+  module_type: "DGLOGBDataModule"
+  args:
+    cache_data_path: null
+    dataset_name: "ogbg-lsc-pcqm4mv2"
+
+    # Weights
+    weights_type: null
+    sample_size: null
+
+    # Featurization
+    featurization_n_jobs: -1
+    featurization_progress: True
+    featurization:
+      atom_property_list_onehot: [atomic-number, valence]
+      atom_property_list_float: [mass, electronegativity, in-ring, hybridization, chirality, aromatic, degree, formal-charge, single-bond, double-bond, radical-electron, vdw-radius, covalent-radius, metal]
+      edge_property_list: [bond-type-onehot, bond-type-float, stereo, in-ring, conjugated, estimated-bond-length]
+      add_self_loop: False
+      explicit_H: False
+      use_bonds_weights: False
+      pos_encoding_as_features: &pos_enc
+        pos_type: laplacian_eigvec
+        num_pos: 3
+        normalization: "none"
+        disconnected_comp: True
+      pos_encoding_as_directions: *pos_enc
+      mask_nan: raise
+
+    # Train, val, test parameters
+    batch_size_train_val: 1024
+    batch_size_test: 1024
+
+  # Data loading
+  num_workers: 0
+  pin_memory: False
+  persistent_workers: False  # Keep True on Windows if running multiple workers
+
+
+architecture:
+  model_type: fulldglnetwork
+  pre_nn:   # Set as null to avoid a pre-nn network
+    out_dim: &hidden_dim 420
+    hidden_dims: *hidden_dim
+    depth: 3
+    activation: relu
+    last_activation: none
+    dropout: &dropout_mlp 0.2
+    last_dropout: *dropout_mlp
+    normalization: &normalization "batch_norm"
+    last_normalization: *normalization
+    residual_type: simple
+
+  pre_nn_edges:   # Set as null to avoid a pre-nn network
+    out_dim: 32
+    hidden_dims: 32
+    depth: 3
+    activation: relu
+    last_activation: none
+    dropout: *dropout_mlp
+    last_dropout: *dropout_mlp
+    normalization: *normalization
+    last_normalization: *normalization
+    residual_type: simple
+
+  gnn:  # Set as null to avoid a post-nn network
+    out_dim: *hidden_dim
+    hidden_dims: *hidden_dim
+    depth: 16
+    activation: none
+    last_activation: none
+    dropout: &dropout_gnn 0.2
+    last_dropout: *dropout_gnn
+    normalization: *normalization
+    last_normalization: *normalization
+    residual_type: simple
+    pooling: ['sum', 'max', 'dir1']
+    virtual_node: 'none'
+    layer_type: 'dgn-msgpass'
+    layer_kwargs:
+      # num_heads: 3
+      aggregators: [mean, max, sum, dir1/dx_abs]
+      scalers: [identity]
+
+  post_nn:
+    out_dim: 1
+    hidden_dims: *hidden_dim
+    depth: 3
+    activation: relu
+    last_activation: none
+    dropout: *dropout_mlp
+    last_dropout: 0.
+    normalization: *normalization
+    last_normalization: "none"
+    residual_type: simple
+
+predictor:
+  metrics_on_progress_bar: ["mae", "pearsonr"]
+  loss_fun: mse
+  random_seed: *seed
+  optim_kwargs:
+    lr: 5.e-3
+    weight_decay: 0
+  lr_reduce_on_plateau_kwargs:
+    factor: 0.5
+    patience: 20
+    min_lr: 2.e-4
+  scheduler_kwargs:
+    monitor: &monitor mae/val
+    mode: &mode min
+    frequency: 1
+  target_nan_mask: ignore-flatten # null: no mask, 0: 0 mask, ignore: ignore nan values from loss
+
+
+metrics:
+  - name: mae
+    metric: mae
+    threshold_kwargs: null
+    target_nan_mask: ignore-flatten
+
+
+  - name: pearsonr
+    metric: pearsonr
+    threshold_kwargs: null
+    target_nan_mask: ignore-flatten
+
+  - name: mse
+    metric: mse
+    threshold_kwargs: null
+    target_nan_mask: ignore-flatten
+
+  - name: spearmanr
+    metric: spearmanr
+    threshold_kwargs: null
+    target_nan_mask: ignore-flatten
+
+
+  - name: f1 > 5
+    metric: f1
+    num_classes: 2
+    average: micro
+    threshold_kwargs: &threshold_1
+      operator: greater
+      threshold: 5
+      th_on_preds: True
+      th_on_target: True
+      target_to_int: True
+
+  - name: f1 > 4
+    metric: f1
+    num_classes: 2
+    average: micro
+    threshold_kwargs: &threshold_2
+      operator: greater
+      threshold: 4
+      th_on_preds: True
+      th_on_target: True
+      target_to_int: True
+
+
+trainer:
+  logger:
+    save_dir: logs/ogb-molpcqm4mv2
+  early_stopping:
+    monitor: *monitor
+    min_delta: 0
+    patience: 80
+    mode: *mode
+  model_checkpoint:
+    dirpath: models_checkpoints/ogb-molpcqm4mv2/
+    filename: "model"
+    monitor: *monitor
+    mode: *mode
+    save_top_k: 1
+    every_n_val_epochs: 1
+  trainer:
+    max_epochs: 1000
+    min_epochs: 100
+    gpus: 1
+    accumulate_grad_batches: 1
+
diff --git a/expts/main_run.py b/expts/main_run.py
index 916dad81c..2bb1467c9 100644
--- a/expts/main_run.py
+++ b/expts/main_run.py
@@ -16,12 +16,10 @@
 # CONFIG_FILE = "expts/config_molPCBA.yaml"
 # CONFIG_FILE = "expts/config_molHIV.yaml"
 # CONFIG_FILE = "expts/config_molPCQM4M.yaml"
+CONFIG_FILE = "expts/config_molPCQM4Mv2.yaml"
 # CONFIG_FILE = "expts/config_micro_ZINC.yaml"
 # CONFIG_FILE = "expts/config_micro-PCBA.yaml"
 # CONFIG_FILE = "expts/config_ZINC_bench_gnn.yaml"
-# CONFIG_FILE = "expts_untracked/config_htsfp_pcba.yaml"
-# CONFIG_FILE = "expts_untracked/config_mega_pubchem.yaml"
-CONFIG_FILE = "expts_untracked/config_mega_pubchem_L1000.yaml"
 os.chdir(MAIN_DIR)
 
 
diff --git a/goli/config/_loader.py b/goli/config/_loader.py
index 55b9e4609..2dc38421f 100644
--- a/goli/config/_loader.py
+++ b/goli/config/_loader.py
@@ -10,18 +10,12 @@
 
 from goli.trainer.metrics import MetricWrapper
 from goli.nn import FullDGLNetwork, FullDGLSiameseNetwork, FeedForwardNN
-from goli.data.datamodule import DGLFromSmilesDataModule, DGLOGBDataModule
 from goli.trainer.predictor import PredictorModule
+from goli.utils.spaces import DATAMODULE_DICT
 
 from pytorch_lightning.loggers.tensorboard import TensorBoardLogger
 
 
-DATAMODULE_DICT = {
-    "DGLFromSmilesDataModule": DGLFromSmilesDataModule,
-    "DGLOGBDataModule": DGLOGBDataModule,
-}
-
-
 def load_datamodule(
     config: Union[omegaconf.DictConfig, Dict[str, Any]],
 ):
diff --git a/goli/data/collate.py b/goli/data/collate.py
index 84245a44a..33ebd2caa 100644
--- a/goli/data/collate.py
+++ b/goli/data/collate.py
@@ -2,14 +2,13 @@
 from torch.utils.data.dataloader import default_collate
 import dgl
 from inspect import signature, _empty
-from typing import Union
+from typing import Union, List
 
 from goli.features import dgl_dict_to_graph
 
 
 def goli_collate_fn(
-    elements,
-    mask_nan: Union[str, float, type(None)] = "raise",
+    elements, mask_nan: Union[str, float, type(None)] = "raise", do_not_collate_keys: List[str] = []
 ):
     """This collate function is identical to the default
     pytorch collate function but add support for `dgl.DGLGraph`
@@ -41,6 +40,9 @@ def goli_collate_fn(
             - "warn": Raise a warning when there is a nan or inf in the featurization
             - "None": DEFAULT. Don't do anything
             - "Floating value": Replace nans or inf by the specified value
+
+        do_not_batch_keys:
+            Keys to ignore for the collate
     """
 
     elem = elements[0]
@@ -64,6 +66,10 @@ def goli_collate_fn(
             elif isinstance(elem[key], dgl.DGLGraph):
                 batch[key] = dgl.batch([d[key] for d in elements])
 
+            # Ignore the collate for specific keys
+            elif key in do_not_collate_keys:
+                batch[key] = [d[key] for d in elements]
+
             # Otherwise, use the default torch batching
             else:
                 batch[key] = default_collate([d[key] for d in elements])
diff --git a/goli/data/datamodule.py b/goli/data/datamodule.py
index 0b092f135..bee93ed06 100644
--- a/goli/data/datamodule.py
+++ b/goli/data/datamodule.py
@@ -4,6 +4,7 @@
 from functools import partial
 import importlib.resources
 import zipfile
+from copy import deepcopy
 
 from loguru import logger
 import fsspec
@@ -29,6 +30,32 @@
 from torch.utils.data.dataloader import DataLoader, Dataset
 from torch.utils.data import Subset
 
+PCQM4M_meta = {
+    "num tasks": 1,
+    "eval metric": "mae",
+    "download_name": "pcqm4m_kddcup2021",
+    "url": "https://dgl-data.s3-accelerate.amazonaws.com/dataset/OGB-LSC/pcqm4m_kddcup2021.zip",
+    "data type": "mol",
+    "has_node_attr": True,
+    "has_edge_attr": True,
+    "task type": "regression",
+    "num classes": -1,
+    "split": "scaffold",
+    "additional node files": "None",
+    "additional edge files": "None",
+    "binary": False,
+    "version": 1,
+}
+
+PCQM4Mv2_meta = deepcopy(PCQM4M_meta)
+PCQM4Mv2_meta.update(
+    {
+        "download_name": "pcqm4m-v2",
+        "url": "https://dgl-data.s3-accelerate.amazonaws.com/dataset/OGB-LSC/pcqm4m-v2.zip",
+        "version": 2,
+    }
+)
+
 
 class DGLDataset(Dataset):
     def __init__(
@@ -95,6 +122,7 @@ def __init__(
         self.train_ds = None
         self.val_ds = None
         self.test_ds = None
+        self._predict_ds = None
 
     def prepare_data(self):
         raise NotImplementedError()
@@ -127,7 +155,7 @@ def test_dataloader(self, **kwargs):
     def predict_dataloader(self, **kwargs):
 
         return self._dataloader(
-            dataset=self.dataset,  # type: ignore
+            dataset=self.predict_ds,  # type: ignore
             batch_size=self.batch_size_test,
             shuffle=False,
         )
@@ -148,6 +176,19 @@ def num_node_feats(self):
     def num_edge_feats(self):
         raise NotImplementedError()
 
+    @property
+    def predict_ds(self):
+        """Get the dataset used for the prediction"""
+        if self._predict_ds is None:
+            return self.test_ds
+        else:
+            return self._predict_ds
+
+    @predict_ds.setter
+    def predict_ds(self, value):
+        """Set the dataset for the prediction"""
+        self._predict_ds = value
+
     def get_first_graph(self):
         raise NotImplementedError()
 
@@ -224,6 +265,7 @@ def __init__(
         featurization_backend: str = "multiprocessing",
         collate_fn: Optional[Callable] = None,
         prepare_dict_or_graph: str = "dict",
+        dataset_class: type = DGLDataset,
     ):
         """
 
@@ -296,6 +338,7 @@ def __init__(
                 - "dict": Process molecules as a Dict. It's faster and requires less RAM during
                   pre-processing, but slower during training since DGLGraphs will be created
                   during data-loading.
+            dataset_class: The class used to create the dataset from which to sample.
         """
         super().__init__(
             batch_size_train_val=batch_size_train_val,
@@ -335,9 +378,11 @@ def __init__(
         self.train_ds = None
         self.val_ds = None
         self.test_ds = None
+        self._predict_ds = None
         self.train_indices = None
         self.val_indices = None
         self.test_indices = None
+        self.dataset_class = dataset_class
 
         if prepare_dict_or_graph == "dict":
             self.smiles_transformer = partial(mol_to_dglgraph_dict, **featurization)
@@ -384,7 +429,7 @@ def prepare_data(self):
         logger.info(f"Prepare dataset with {len(df)} data points.")
 
         # Extract smiles and labels
-        smiles, labels, indices, weights, sample_idx = self._extract_smiles_labels(
+        smiles, labels, sample_idx, extras = self._extract_smiles_labels(
             df,
             smiles_col=self.smiles_col,
             label_cols=self.label_cols,
@@ -393,40 +438,13 @@ def prepare_data(self):
             weights_type=self.weights_type,
         )
 
-        # Precompute the features
-        # NOTE(hadim): in case of very large dataset we could:
-        # - or cache the data and read from it during `next(iter(dataloader))`
-        # - or compute the features on-the-fly during `next(iter(dataloader))`
-        # For now we compute in advance and hold everything in memory.
-
-        if self.featurization_n_jobs == 0:
-            features = [self.smiles_transformer(s) for s in tqdm(smiles)]
-        else:
-            features = Parallel(n_jobs=self.featurization_n_jobs, backend=self.featurization_backend)(
-                delayed(self.smiles_transformer)(s) for s in tqdm(smiles)
-            )
+        # Convert SMILES to features (graphs, fingerprints, etc.)
+        features, idx_none = self._featurize_molecules(smiles, sample_idx)
 
-        # Warn about None molecules
-        is_none = np.array([ii for ii, feat in enumerate(features) if feat is None])
-        if len(is_none) > 0:
-            mols_to_msg = [f"{sample_idx[idx]}: {smiles[idx]}" for idx in is_none]
-            msg = "\n".join(mols_to_msg)
-            logger.warning(
-                (f"{len(is_none)} molecules will be removed since they failed featurization:\n" + msg)
-            )
-
-        # Remove None molecules
-        if len(is_none) > 0:
-            df.drop(df.index[is_none], axis=0)
-            features = [feat for feat in features if not (feat is None)]
-            sample_idx = np.delete(sample_idx, is_none, axis=0)
-            smiles = np.delete(smiles, is_none, axis=0)
-            if labels is not None:
-                labels = np.delete(labels, is_none, axis=0)
-            if weights is not None:
-                weights = np.delete(weights, is_none, axis=0)
-            if indices is not None:
-                indices = np.delete(indices, is_none, axis=0)
+        # Filter the molecules, labels, etc. for the molecules that failed featurization
+        df, features, smiles, labels, sample_idx, extras = self._filter_none_molecules(
+            idx_none, df, features, smiles, labels, sample_idx, extras
+        )
 
         # Get splits indices
         self.train_indices, self.val_indices, self.test_indices = self._get_split_indices(
@@ -439,12 +457,11 @@ def prepare_data(self):
         )
 
         # Make the torch datasets (mostly a wrapper there is no memory overhead here)
-        self.dataset = DGLDataset(
+        self.dataset = self.dataset_class(
             smiles=smiles,
             features=features,
             labels=labels,
-            indices=indices,
-            weights=weights,
+            **extras,
         )
 
         self._save_to_cache()
@@ -459,6 +476,101 @@ def setup(self, stage: str = None):
         if stage == "test" or stage is None:
             self.test_ds = Subset(self.dataset, self.test_indices)  # type: ignore
 
+    def _featurize_molecules(
+        self, smiles: Iterable[str], sample_idx: Iterable[int] = None
+    ) -> Tuple[List, List]:
+        """
+        Precompute the features (graphs, fingerprints, etc.) from the SMILES.
+        Features are computed from `self.smiles_transformer`.
+        A warning is issued to mention which molecules failed featurization.
+
+        Note:
+            (hadim): in case of very large dataset we could:
+            - or cache the data and read from it during `next(iter(dataloader))`
+            - or compute the features on-the-fly during `next(iter(dataloader))`
+            For now we compute in advance and hold everything in memory.
+
+        Parameters:
+            smiles: A list of all the molecular SMILES to featurize
+            sample_idx: The indexes corresponding to the sampled SMILES.
+                If not provided, computed from `numpy.arange`.
+
+        Returns:
+            features: A list of all the featurized molecules
+            idx_none: A list of the indexes that failed featurization
+        """
+
+        # Loop all the smiles and compute the features
+        if self.featurization_n_jobs == 0:
+            features = [self.smiles_transformer(s) for s in tqdm(smiles)]
+        else:
+            features = Parallel(n_jobs=self.featurization_n_jobs, backend=self.featurization_backend)(
+                delayed(self.smiles_transformer)(s) for s in tqdm(smiles)
+            )
+
+        # Warn about None molecules
+        if sample_idx is None:
+            sample_idx = np.arange(len(smiles))
+        idx_none = [ii for ii, feat in enumerate(features) if feat is None]
+        if len(idx_none) > 0:
+            mols_to_msg = [f"{sample_idx[idx]}: {smiles[idx]}" for idx in idx_none]
+            msg = "\n".join(mols_to_msg)
+            logger.warning(
+                (f"{len(idx_none)} molecules will be removed since they failed featurization:\n" + msg)
+            )
+
+        return features, idx_none
+
+    @staticmethod
+    def _filter_none_molecules(
+        idx_none: Iterable,
+        *args: Union[pd.DataFrame, pd.Series, np.ndarray, torch.Tensor, list, tuple, Dict[Any, Iterable]],
+    ) -> List[Union[pd.DataFrame, pd.Series, np.ndarray, torch.Tensor, list, tuple, Dict[Any, Iterable]]]:
+        """
+        Filter the molecules, labels, etc. for the molecules that failed featurization.
+
+        Parameters:
+            idx_none: A list of the indexes that failed featurization
+            args: Any argument from which to filter the failed SMILES.
+                Can be a `list`, `tuple`, `Tensor`, `np.array`, `Dict`, `pd.DataFrame`, `pd.Series`.
+                Otherwise, it is not filtered.
+                WARNING: If a `pd.DataFrame` or `pd.Series` is passed, it filters by the row indexes,
+                NOT by the `DataFrame.index` or `Series.index`! Be careful!
+
+        Returns:
+            out: All the `args` with the indexes from `idx_none` removed.
+        """
+        if len(idx_none) == 0:
+            return args
+        idx_none = np.asarray(idx_none)
+
+        out = []
+        for arg in args:
+            if isinstance(arg, pd.DataFrame):
+                new = arg.drop(arg.index[idx_none], axis=0)
+            elif isinstance(arg, pd.Series):
+                new = arg.drop(arg.index[idx_none], axis=0)
+            elif isinstance(arg, np.ndarray):
+                new = np.delete(arg, idx_none, axis=0)
+            elif isinstance(arg, torch.Tensor):
+                not_none = torch.ones(arg.shape[0], dtype=bool)
+                not_none[idx_none] = False
+                new = arg[not_none]
+            elif isinstance(arg, (list, tuple)):
+                arg = list(arg)
+                new = [elem for ii, elem in enumerate(arg) if ii not in idx_none]
+            elif isinstance(arg, dict):
+                new = {}
+                for key, val in arg.items():
+                    new[key] = DGLFromSmilesDataModule._filter_none_molecules(idx_none, val)
+            else:
+                new = arg
+            out.append(new)
+
+        out = tuple(out) if len(out) > 1 else out[0]
+
+        return out
+
     def _parse_label_cols(self, label_cols: Union[type(None), str, List[str]], smiles_col: str) -> List[str]:
         r"""
         Parse the choice of label columns depending on the type of input.
@@ -496,9 +608,9 @@ def is_prepared(self):
 
     @property
     def is_setup(self):
-        if not hasattr(self, "train_ds"):
+        if not (hasattr(self, "train_ds") or hasattr(self, "test_ds")):
             return False
-        return getattr(self, "train_ds") is not None
+        return (getattr(self, "train_ds") is not None) or (getattr(self, "test_ds") is not None)
 
     @property
     def num_node_feats(self):
@@ -547,7 +659,7 @@ def get_first_graph(self):
         else:
             df = self.df.iloc[0:10, :]
 
-        smiles, _, _, _, _ = self._extract_smiles_labels(df, self.smiles_col, self.label_cols)
+        smiles, _, _, _ = self._extract_smiles_labels(df, self.smiles_col, self.label_cols)
 
         graph = None
         for s in smiles:
@@ -655,7 +767,7 @@ def _extract_smiles_labels(
         weights_col: str = None,
         weights_type: str = None,
     ) -> Tuple[
-        np.ndarray, np.ndarray, Union[type(None), np.ndarray], Union[type(None), np.ndarray], np.ndarray
+        np.ndarray, np.ndarray, Union[type(None), np.ndarray], Dict[str, Union[type(None), np.ndarray]]
     ]:
         """For a given dataframe extract the SMILES and labels columns. Smiles is returned as a list
         of string while labels are returned as a 2D numpy array.
@@ -716,7 +828,8 @@ def _extract_smiles_labels(
 
             weights /= np.max(weights)  # Put the max weight to 1
 
-        return smiles, labels, indices, weights, sample_idx
+        extras = {"indices": indices, "weights": weights}
+        return smiles, labels, sample_idx, extras
 
     def _get_split_indices(
         self,
@@ -912,10 +1025,7 @@ def _load_dataset(self, metadata: dict):
         """Download, extract and load an OGB dataset."""
 
         base_dir = fs.get_cache_dir("ogb")
-        if metadata["download_name"] == "pcqm4m":
-            dataset_dir = base_dir / (metadata["download_name"] + "_kddcup2021")
-        else:
-            dataset_dir = base_dir / metadata["download_name"]
+        dataset_dir = base_dir / metadata["download_name"]
 
         if not dataset_dir.exists():
 
@@ -932,7 +1042,7 @@ def _load_dataset(self, metadata: dict):
             zf.extractall(base_dir)
 
         # Load CSV file
-        if metadata["download_name"] == "pcqm4m":
+        if metadata["download_name"].startswith("pcqm4m"):
             df_path = dataset_dir / "raw" / "data.csv.gz"
         else:
             df_path = dataset_dir / "mapping" / "mol.csv.gz"
@@ -940,12 +1050,16 @@ def _load_dataset(self, metadata: dict):
         df = pd.read_csv(df_path)
 
         # Load split from the OGB dataset and save them in a single CSV file
-        if metadata["download_name"] == "pcqm4m":
+        if metadata["download_name"].startswith("pcqm4m"):
             split_name = metadata["split"]
             split_dict = torch.load(dataset_dir / "split_dict.pt")
             train_split = pd.DataFrame(split_dict["train"])
             val_split = pd.DataFrame(split_dict["valid"])
-            test_split = pd.DataFrame(split_dict["test"])
+            if "test" in split_dict.keys():
+                test_split = pd.DataFrame(split_dict["test"])
+            else:
+                test_split = pd.DataFrame(split_dict["test-dev"])
+
             splits = pd.concat([train_split, val_split, test_split], axis=1)  # type: ignore
             splits.columns = ["train", "val", "test"]
 
@@ -971,7 +1085,7 @@ def _load_dataset(self, metadata: dict):
             splits.to_csv(splits_path, index=None)
 
         # Get column names: OGB columns are predictable
-        if metadata["download_name"] == "pcqm4m":
+        if metadata["download_name"].startswith("pcqm4m"):
             idx_col = df.columns[0]
             smiles_col = df.columns[-2]
             label_cols = df.columns[-1:].to_list()
@@ -984,7 +1098,6 @@ def _load_dataset(self, metadata: dict):
 
     def _get_dataset_metadata(self, dataset_name: str):
         ogb_metadata = self._get_ogb_metadata()
-
         if dataset_name not in ogb_metadata.index:
             raise ValueError(f"'{dataset_name}' is not a valid dataset name.")
 
@@ -995,10 +1108,13 @@ def _get_ogb_metadata(self):
 
         with importlib.resources.open_text("ogb.graphproppred", "master.csv") as f:
             ogb_metadata = pd.read_csv(f)
-
         ogb_metadata = ogb_metadata.set_index(ogb_metadata.columns[0])
         ogb_metadata = ogb_metadata.T
 
+        # Add metadata related to PCQM4M
+        ogb_metadata = ogb_metadata.append(pd.DataFrame(PCQM4M_meta, index=["ogbg-lsc-pcqm4m"]))
+        ogb_metadata = ogb_metadata.append(pd.DataFrame(PCQM4Mv2_meta, index=["ogbg-lsc-pcqm4mv2"]))
+
         # Only keep datasets of type 'mol'
         ogb_metadata = ogb_metadata[ogb_metadata["data type"] == "mol"]
 
diff --git a/goli/nn/architectures.py b/goli/nn/architectures.py
index cfbf07736..6f1e25d73 100644
--- a/goli/nn/architectures.py
+++ b/goli/nn/architectures.py
@@ -12,7 +12,6 @@
     ResidualConnectionRandom,
 )
 from goli.nn.dgl_layers.pooling import parse_pooling_layer, VirtualNode
-from goli.utils.spaces import LAYERS_DICT, RESIDUALS_DICT
 
 
 class FeedForwardNN(nn.Module):
@@ -135,6 +134,8 @@ def __init__(
         self.name = name
 
         # Parse the layer and residuals
+        from goli.utils.spaces import LAYERS_DICT, RESIDUALS_DICT
+
         self.layer_class, self.layer_name = self._parse_class_from_dict(layer_type, LAYERS_DICT)
         self.residual_class, self.residual_name = self._parse_class_from_dict(residual_type, RESIDUALS_DICT)
 
diff --git a/goli/trainer/metrics.py b/goli/trainer/metrics.py
index dfb4f0a67..a5e38dc61 100644
--- a/goli/trainer/metrics.py
+++ b/goli/trainer/metrics.py
@@ -1,26 +1,7 @@
 from typing import Union, Callable, Optional, Dict, Any
-
-from copy import deepcopy
-
 import torch
-
 import operator as op
-
 from torchmetrics.utilities import reduce
-from torchmetrics.functional import auroc
-from torchmetrics.functional import (
-    accuracy,
-    average_precision,
-    confusion_matrix,
-    f1,
-    fbeta,
-    precision_recall_curve,
-    precision,
-    recall,
-    auroc,
-    mean_absolute_error,
-    mean_squared_error,
-)
 
 from goli.utils.tensor import nan_mean
 
@@ -170,29 +151,6 @@ def spearmanr(preds: torch.Tensor, target: torch.Tensor, reduction: str = "eleme
     return spearman
 
 
-METRICS_CLASSIFICATION = {
-    "accuracy": accuracy,
-    "averageprecision": average_precision,
-    "auroc": auroc,
-    "confusionmatrix": confusion_matrix,
-    "f1": f1,
-    "fbeta": fbeta,
-    "precisionrecallcurve": precision_recall_curve,
-    "precision": precision,
-    "recall": recall,
-}
-
-METRICS_REGRESSION = {
-    "mae": mean_absolute_error,
-    "mse": mean_squared_error,
-    "pearsonr": pearsonr,
-    "spearmanr": spearmanr,
-}
-
-METRICS_DICT = deepcopy(METRICS_CLASSIFICATION)
-METRICS_DICT.update(METRICS_REGRESSION)
-
-
 class MetricWrapper:
     r"""
     Allows to initialize a metric from a name or Callable, and initialize the
@@ -233,6 +191,7 @@ def __init__(
             kwargs:
                 Other arguments to call with the metric
         """
+        from goli.utils.spaces import METRICS_DICT
 
         self.metric = METRICS_DICT[metric] if isinstance(metric, str) else metric
 
diff --git a/goli/trainer/predictor.py b/goli/trainer/predictor.py
index 12fbbad28..2cc74dc30 100644
--- a/goli/trainer/predictor.py
+++ b/goli/trainer/predictor.py
@@ -5,10 +5,11 @@
 from copy import deepcopy
 import yaml
 import dgl
+from loguru import logger
+from inspect import signature
 
 import torch
 from torch import nn, Tensor
-from torch.optim.lr_scheduler import ReduceLROnPlateau
 
 import pytorch_lightning as pl
 from pytorch_lightning import _logger as log
@@ -17,13 +18,7 @@
 from goli.config.config_convert import recursive_config_reformating
 from goli.utils.tensor import nan_mean, nan_std, nan_median
 from goli.utils.fs import mkdir
-
-LOSS_DICT = {
-    "mse": torch.nn.MSELoss(),
-    "bce": torch.nn.BCELoss(),
-    "l1": torch.nn.L1Loss(),
-    "mae": torch.nn.L1Loss(),
-}
+from goli.utils.spaces import LOSS_DICT, SCHEDULER_DICT
 
 GOLI_PRETRAINED_MODELS = {
     "goli-zinc-micro-dummy-test": "gcs://goli-public/pretrained-models/goli-zinc-micro-dummy-test/model.ckpt"
@@ -127,13 +122,13 @@ def __init__(
         random_seed: int = 42,
         optim_kwargs: Optional[Dict[str, Any]] = None,
         lr_reduce_on_plateau_kwargs: Optional[Dict[str, Any]] = None,
+        torch_scheduler_kwargs: Optional[Dict[str, Any]] = None,
         scheduler_kwargs: Optional[Dict[str, Any]] = None,
         target_nan_mask: Optional[Union[int, float, str]] = None,
         metrics: Dict[str, Callable] = None,
         metrics_on_progress_bar: List[str] = [],
         metrics_on_training_set: Optional[List[str]] = None,
-        flag_n_steps: int = 0,
-        flag_alpha: float = 0.01,
+        flag_kwargs: Dict[str, Any] = None,
     ):
         r"""
         A class that allows to use regression or classification models easily
@@ -161,6 +156,7 @@ def __init__(
                 - weight_decay `float`: Weight decay used to regularize the optimizer (Default=`0.`)
 
             lr_reduce_on_plateau_kwargs:
+                DEPRECIATED. USE `torch_scheduler_kwargs` instead.
                 Dictionnary for the reduction of learning rate when reaching plateau, with possible keys below.
 
                 - factor `float`: Factor by which to reduce the learning rate (Default=`0.5`)
@@ -172,8 +168,15 @@ def __init__(
                 - min_lr `float`: A scalar or a list of scalars. A lower bound on the learning rate
                   of all param groups or each group respectively (Default=`1e-4`)
 
+            torch_scheduler_kwargs:
+                Dictionnary for the scheduling of learning rate, with possible keys below.
+
+                - type `str`: Type of the learning rate to use from pytorch. Examples are
+                  `'ReduceLROnPlateau'` (default), `'CosineAnnealingWarmRestarts'`, `'StepLR'`, etc.
+                - **kwargs: Any other argument for the learning rate scheduler
+
             scheduler_kwargs:
-                Dictionnary for the scheduling of the learning rate modification
+                Dictionnary for the scheduling of the learning rate modification used by pytorch-lightning
 
                 - monitor `str`: metric to track (Default=`"loss/val"`)
                 - interval `str`: Whether to look at iterations or epochs (Default=`"epoch"`)
@@ -183,14 +186,18 @@ def __init__(
                 - frequency `int`: **TODO: NOT REALLY SURE HOW IT WORKS!** (Default=`1`)
 
             target_nan_mask:
-                TODO: It's not implemented for the metrics yet!!
 
-                - None: Do not change behaviour if there are nans
+                - None: Do not change behaviour if there are NaNs
+
+                - int, float: Value used to replace NaNs. For example, if `target_nan_mask==0`, then
+                  all NaNs will be replaced by zeros
 
-                - int, float: Value used to replace nans. For example, if `target_nan_mask==0`, then
-                  all nans will be replaced by zeros
+                - 'ignore-flatten': The Tensor will be reduced to a vector without the NaN values.
 
-                - 'ignore': Nans will be ignored when computing the loss.
+                - 'ignore-mean-label': NaNs will be ignored when computing the loss. Note that each column
+                  has a different number of NaNs, so the metric will be computed separately
+                  on each column, and the metric result will be averaged over all columns.
+                  *This option might slowdown the computation if there are too many labels*
 
             metrics:
                 A dictionnary of metrics to compute on the prediction, other than the loss function.
@@ -204,13 +211,15 @@ def __init__(
                 If `None`, all the metrics are computed. Using less metrics can significantly improve
                 performance, depending on the number of readouts.
 
-            flag_n_steps:
-                An integer that specifies the number of ascent steps when running FLAG during training.
-                Default value of 0 trains GNNs without FLAG, and any value greater than 0 will use FLAG with that
-                many iterations.
+            flag_kwargs:
+                Keyword arguments used for FLAG, and adversarial data augmentation for graph networks.
+                See: https://arxiv.org/abs/2010.09891
+
+                - n_steps: An integer that specifies the number of ascent steps when running FLAG during training.
+                  Default value of 0 trains GNNs without FLAG, and any value greater than 0 will use FLAG with that
+                  many iterations.
 
-            flag_alpha:
-                A float that specifies the ascent step size when running FLAG.
+                - alpha: A float that specifies the ascent step size when running FLAG. Default=0.01
         """
 
         self.save_hyperparameters()
@@ -231,33 +240,56 @@ def __init__(
             list(self.metrics.keys()) if metrics_on_training_set is None else metrics_on_training_set
         )
         self.n_params = sum(p.numel() for p in self.parameters() if p.requires_grad)
-        self.lr_reduce_on_plateau_kwargs = lr_reduce_on_plateau_kwargs
-        self.optim_kwargs = optim_kwargs
-        self.scheduler_kwargs = scheduler_kwargs
-        self.flag_n_steps = flag_n_steps
-        self.flag_alpha = flag_alpha
-
-        # Set the default value for the optimizer
-        self.optim_kwargs = optim_kwargs if optim_kwargs is not None else {}
-        self.optim_kwargs.setdefault("lr", 1e-3)
-        self.optim_kwargs.setdefault("weight_decay", 0.0)
-
-        self.lr_reduce_on_plateau_kwargs = (
-            lr_reduce_on_plateau_kwargs if lr_reduce_on_plateau_kwargs is not None else {}
-        )
-        self.lr_reduce_on_plateau_kwargs.setdefault("factor", 0.5)
-        self.lr_reduce_on_plateau_kwargs.setdefault("patience", 10)
-        self.lr_reduce_on_plateau_kwargs.setdefault("min_lr", 1e-4)
-
-        self.optim_kwargs = optim_kwargs if optim_kwargs is not None else {}
-        self.scheduler_kwargs.setdefault("monitor", "loss/val")
-        self.scheduler_kwargs.setdefault("interval", "epoch")
-        self.scheduler_kwargs.setdefault("mode", "min")
-        self.scheduler_kwargs.setdefault("frequency", 1)
-        self.scheduler_kwargs.setdefault("strict", True)
-
-        monitor = scheduler_kwargs["monitor"].split("/")[0]
-        mode = scheduler_kwargs["mode"]
+
+        # Set the parameters and default values for the FLAG adversarial augmentation, and check values
+        self.flag_kwargs = {"alpha": 0.01, "n_steps": 0}
+        if flag_kwargs is not None:
+            self.flag_kwargs.update(flag_kwargs)
+        assert isinstance(self.flag_kwargs["n_steps"], int) and (self.flag_kwargs["n_steps"] >= 0)
+        assert self.flag_kwargs["alpha"] > 0
+
+        # Set the parameters and default value for the optimizer, and check values
+        self.optim_kwargs = {"lr": 1e-3, "weight_decay": 0.0}
+        if optim_kwargs is not None:
+            self.optim_kwargs.update(optim_kwargs)
+        assert self.optim_kwargs["lr"] > 0
+        assert self.optim_kwargs["weight_decay"] >= 0
+
+        # Set the lightning scheduler
+        self.scheduler_kwargs = {
+            "interval": "epoch",
+            "monitor": "loss/val",
+            "mode": "min",
+            "frequency": 1,
+            "strict": True,
+        }
+        if scheduler_kwargs is not None:
+            self.scheduler_kwargs.update(scheduler_kwargs)
+
+        # Set the depreciated arguments, if provided
+        if lr_reduce_on_plateau_kwargs is not None:
+            logger.warning(
+                "`lr_reduce_on_plateau_kwargs` is depreciated, use `torch_scheduler_kwargs` instead."
+            )
+            if torch_scheduler_kwargs is not None:
+                raise ValueError(
+                    "Ambiguous. Both `lr_reduce_on_plateau_kwargs` and `torch_scheduler_kwargs` are provided"
+                )
+            torch_scheduler_kwargs = {
+                "module_type": "ReduceLROnPlateau",
+                **lr_reduce_on_plateau_kwargs,
+            }
+
+        # Set the pytorch scheduler arguments
+        if torch_scheduler_kwargs is None:
+            self.torch_scheduler_kwargs = {}
+        else:
+            self.torch_scheduler_kwargs = torch_scheduler_kwargs
+        self.torch_scheduler_kwargs.setdefault("module_type", "ReduceLROnPlateau")
+
+        # Initialize the epoch summary
+        monitor = self.scheduler_kwargs["monitor"].split("/")[0]
+        mode = self.scheduler_kwargs["mode"]
         self.epoch_summary = EpochSummary(
             monitor, mode=mode, metrics_on_progress_bar=self.metrics_on_progress_bar
         )
@@ -289,13 +321,28 @@ def parse_loss_fun(loss_fun: Union[str, Callable]) -> Callable:
 
         return loss_fun
 
-    def forward(self, inputs: Dict):
+    def forward(self, inputs: Dict) -> Dict[str, Union[torch.Tensor, Any]]:
         r"""
         Returns the result of `self.model.forward(*inputs)` on the inputs.
+        If the output of `out = self.model.forward` is a dictionary with a `"preds"` key,
+        it is returned directly. Otherwise, a new dictionary is created and
+        returns `{"preds": out}`.
+
+        Returns:
+            A dict with a key `"preds"` representing the prediction of the network.
+
         """
+        # Convert to the right dtype and run the model
         feats = self._convert_features_dtype(inputs["features"])
         out = self.model.forward(feats)
-        return out
+
+        # Convert the output of the model to a dictionary
+        if isinstance(out, dict) and ("preds" in out.keys()):
+            out_dict = out
+        else:
+            out_dict = {"preds": out}
+
+        return out_dict
 
     def _convert_features_dtype(self, feats):
         # Convert features to dtype
@@ -312,12 +359,21 @@ def _convert_features_dtype(self, feats):
         return feats
 
     def configure_optimizers(self):
+        # Configure the parameters for the schedulers
+        sc_kwargs = deepcopy(self.torch_scheduler_kwargs)
+        scheduler_class = SCHEDULER_DICT[sc_kwargs.pop("module_type")]
+        sig = signature(scheduler_class.__init__)
+        key_args = [p.name for p in sig.parameters.values() if p.kind == p.KEYWORD_ONLY]
+        if "monitor" in key_args:
+            sc_kwargs.setdefault("monitor", self.scheduler_kwargs["monitor"])
+        if "mode" in key_args:
+            sc_kwargs.setdefault("mode", self.scheduler_kwargs["mode"])
+
+        # Define the optimizer and schedulers
         optimiser = torch.optim.Adam(self.parameters(), **self.optim_kwargs)
-
+        torch_scheduler = scheduler_class(optimizer=optimiser, **sc_kwargs)
         scheduler = {
-            "scheduler": ReduceLROnPlateau(
-                optimizer=optimiser, mode=self.scheduler_kwargs["mode"], **self.lr_reduce_on_plateau_kwargs
-            ),
+            "scheduler": torch_scheduler,
             **self.scheduler_kwargs,
         }
         return [optimiser], [scheduler]
@@ -437,7 +493,7 @@ def _general_step(
         self, batch: Dict[str, Tensor], batch_idx: int, step_name: str, to_cpu: bool
     ) -> Dict[str, Any]:
         r"""Common code for training_step, validation_step and testing_step"""
-        preds = self.forward(batch)
+        preds = self.forward(batch)["preds"]
         targets = batch.pop("labels").to(dtype=preds.dtype)
         weights = batch.pop("weights", None)
 
@@ -468,17 +524,19 @@ def flag_step(
         Github: https://github.com/devnkong/FLAG
         """
 
+        alpha, n_steps = self.flag_kwargs["alpha"], self.flag_kwargs["n_steps"]
+
         X = self._convert_features_dtype(batch["features"])
         X_shape = X.ndata["feat"].shape
 
-        pert = torch.FloatTensor(X_shape).uniform_(-self.flag_alpha, self.flag_alpha).to(device=X.device)
+        pert = torch.FloatTensor(X_shape).uniform_(-alpha, alpha).to(device=X.device)
         pert.requires_grad = True
 
         # Perturb the features
         pert_batch = deepcopy(batch)
         pert_batch["features"].ndata["feat"] = batch["features"].ndata["feat"] + pert
 
-        preds = self.forward(pert_batch)
+        preds = self.forward(pert_batch)["preds"]
         targets = batch.pop("labels").to(dtype=preds.dtype)
         weights = batch.pop("weights", None)
         loss = (
@@ -489,18 +547,18 @@ def flag_step(
                 target_nan_mask=self.target_nan_mask,
                 loss_fun=self.loss_fun,
             )
-            / self.flag_n_steps
+            / n_steps
         )
 
         # Iteratively augment data by applying perturbations
         # Accumulate the gradients to be applied to the weights of the network later on
-        for _ in range(self.flag_n_steps - 1):
+        for _ in range(n_steps - 1):
             loss.backward()
-            pert_data = pert.detach() + self.flag_alpha * torch.sign(pert.grad.detach())
+            pert_data = pert.detach() + alpha * torch.sign(pert.grad.detach())
             pert.data = pert_data.data
             pert.grad[:] = 0
             pert_batch["features"].ndata["feat"] = batch["features"].ndata["feat"] + pert
-            preds = self.forward(pert_batch)
+            preds = self.forward(pert_batch)["preds"]
             loss = (
                 self.compute_loss(
                     preds=preds,
@@ -509,7 +567,7 @@ def flag_step(
                     target_nan_mask=self.target_nan_mask,
                     loss_fun=self.loss_fun,
                 )
-                / self.flag_n_steps
+                / n_steps
             )
 
         device = "cpu" if to_cpu else None
@@ -527,10 +585,10 @@ def training_step(self, batch: Dict[str, Tensor], batch_idx: int) -> Dict[str, A
         step_dict = None
 
         # Train using FLAG
-        if self.flag_n_steps > 0:
+        if self.flag_kwargs["n_steps"] > 0:
             loss, step_dict = self.flag_step(batch=batch, batch_idx=batch_idx, step_name="train", to_cpu=True)
         # Train normally, without using FLAG
-        elif self.flag_n_steps == 0:
+        elif self.flag_kwargs["n_steps"] == 0:
             loss, step_dict = self._general_step(
                 batch=batch, batch_idx=batch_idx, step_name="train", to_cpu=True
             )
diff --git a/goli/utils/spaces.py b/goli/utils/spaces.py
index 8c813b7e4..20a903ea2 100644
--- a/goli/utils/spaces.py
+++ b/goli/utils/spaces.py
@@ -1,6 +1,11 @@
 from copy import deepcopy
+import torch
+import torch.optim.lr_scheduler as sc
+import torchmetrics.functional as met
 
 from goli.nn.base_layers import FCLayer
+from goli.trainer.metrics import pearsonr, spearmanr
+from goli.data.datamodule import DGLFromSmilesDataModule, DGLOGBDataModule
 
 from goli.nn.dgl_layers import (
     GATLayer,
@@ -50,3 +55,49 @@
     "densenet": ResidualConnectionDenseNet,
     "random": ResidualConnectionRandom,
 }
+
+LOSS_DICT = {
+    "mse": torch.nn.MSELoss(),
+    "bce": torch.nn.BCELoss(),
+    "l1": torch.nn.L1Loss(),
+    "mae": torch.nn.L1Loss(),
+}
+
+SCHEDULER_DICT = {
+    "CosineAnnealingLR": sc.CosineAnnealingLR,
+    "CosineAnnealingWarmRestarts": sc.CosineAnnealingWarmRestarts,
+    "CyclicLR": sc.CyclicLR,
+    "ExponentialLR": sc.ExponentialLR,
+    "LambdaLR": sc.LambdaLR,
+    "MultiStepLR": sc.MultiStepLR,
+    "ReduceLROnPlateau": sc.ReduceLROnPlateau,
+    "StepLR": sc.StepLR,
+}
+
+METRICS_CLASSIFICATION = {
+    "accuracy": met.accuracy,
+    "averageprecision": met.average_precision,
+    "auroc": met.auroc,
+    "confusionmatrix": met.confusion_matrix,
+    "f1": met.f1,
+    "fbeta": met.fbeta,
+    "precisionrecallcurve": met.precision_recall_curve,
+    "precision": met.precision,
+    "recall": met.recall,
+}
+
+METRICS_REGRESSION = {
+    "mae": met.mean_absolute_error,
+    "mse": met.mean_squared_error,
+    "pearsonr": pearsonr,
+    "spearmanr": spearmanr,
+}
+
+METRICS_DICT = deepcopy(METRICS_CLASSIFICATION)
+METRICS_DICT.update(METRICS_REGRESSION)
+
+
+DATAMODULE_DICT = {
+    "DGLFromSmilesDataModule": DGLFromSmilesDataModule,
+    "DGLOGBDataModule": DGLOGBDataModule,
+}


Test patch:
diff --git a/tests/test_datamodule.py b/tests/test_datamodule.py
index 3f203af7d..092fca16c 100644
--- a/tests/test_datamodule.py
+++ b/tests/test_datamodule.py
@@ -2,6 +2,9 @@
 import tempfile
 
 import unittest as ut
+import numpy as np
+import torch
+import pandas as pd
 
 import goli
 
@@ -187,6 +190,97 @@ def test_datamodule_cache_invalidation(self):
         assert datam.num_node_feats == 2
         assert datam.num_edge_feats == 8
 
+    def test_none_filtering(self):
+        # Create the objects to filter
+        list_of_num = [ii for ii in range(100)]
+        list_of_str = [str(ii) for ii in list_of_num]
+        tuple_of_num = tuple(list_of_num)
+        array_of_num = np.asarray(list_of_num)
+        array_of_str = np.asarray(list_of_str)
+        tensor_of_num = torch.as_tensor(array_of_num)
+        arrays_of_num = np.stack([list_of_num, list_of_num, list_of_num], axis=1)
+        arrays_of_str = np.stack([list_of_str, list_of_str, list_of_str], axis=1)
+        tensors_of_num = torch.as_tensor(arrays_of_num)
+        dic = {"str": list_of_str, "num": list_of_num}
+        df = pd.DataFrame(dic)
+        df_shuffled = df.sample(frac=1)
+        series_num = df["num"]
+        series_num_shuffled = df_shuffled["num"]
+
+        # Create different indexes to use for filtering
+        all_idx_none = [[3, 17, 88], [22, 33, 44, 55, 66, 77, 88], [], np.arange(len(list_of_num))]
+
+        # Loop all the indexes and filter the objects.
+        for ii, idx_none in enumerate(all_idx_none):
+            msg = f"Failed for ii={ii}"
+
+            # Create the true filtered sequences
+            filtered_num = [ii for ii in range(100) if ii not in idx_none]
+            filtered_str = [str(ii) for ii in filtered_num]
+            assert len(filtered_num) == len(list_of_num) - len(idx_none)
+            assert len(filtered_str) == len(list_of_str) - len(idx_none)
+
+            # Filter the sequences from the Datamodule function
+            (
+                list_of_num_2,
+                list_of_str_2,
+                tuple_of_num_2,
+                array_of_num_2,
+                array_of_str_2,
+                tensor_of_num_2,
+                df_2,
+                df_shuffled_2,
+                dic_2,
+                arrays_of_num_2,
+                arrays_of_str_2,
+                tensors_of_num_2,
+                series_num_2,
+                series_num_shuffled_2,
+            ) = goli.data.DGLFromSmilesDataModule._filter_none_molecules(
+                idx_none,
+                list_of_num,
+                list_of_str,
+                tuple_of_num,
+                array_of_num,
+                array_of_str,
+                tensor_of_num,
+                df,
+                df_shuffled,
+                dic,
+                arrays_of_num,
+                arrays_of_str,
+                tensors_of_num,
+                series_num,
+                series_num_shuffled,
+            )
+
+            df_shuffled_2 = df_shuffled_2.sort_values(by="num", axis=0)
+            series_num_shuffled_2 = series_num_shuffled_2.sort_values(axis=0)
+
+            # Assert the filtering is done correctly
+            self.assertListEqual(list_of_num_2, filtered_num, msg=msg)
+            self.assertListEqual(list_of_str_2, filtered_str, msg=msg)
+            self.assertListEqual(list(tuple_of_num_2), filtered_num, msg=msg)
+            self.assertListEqual(array_of_num_2.tolist(), filtered_num, msg=msg)
+            self.assertListEqual(array_of_str_2.tolist(), filtered_str, msg=msg)
+            self.assertListEqual(tensor_of_num_2.tolist(), filtered_num, msg=msg)
+            for jj in range(arrays_of_num.shape[1]):
+                self.assertListEqual(arrays_of_num_2[:, jj].tolist(), filtered_num, msg=msg)
+                self.assertListEqual(arrays_of_str_2[:, jj].tolist(), filtered_str, msg=msg)
+                self.assertListEqual(tensors_of_num_2[:, jj].tolist(), filtered_num, msg=msg)
+            self.assertListEqual(dic_2["num"], filtered_num, msg=msg)
+            self.assertListEqual(dic_2["str"], filtered_str, msg=msg)
+            self.assertListEqual(df_2["num"].tolist(), filtered_num, msg=msg)
+            self.assertListEqual(df_2["str"].tolist(), filtered_str, msg=msg)
+            self.assertListEqual(series_num_2.tolist(), filtered_num, msg=msg)
+
+            # When the dataframe is shuffled, the lists are different because the filtering
+            # is done on the row indexes, not the dataframe indexes.
+            bool_to_check = (len(idx_none) == 0) or (len(idx_none) == len(df_shuffled))
+            self.assertIs(df_shuffled_2["num"].tolist() == filtered_num, bool_to_check, msg=msg)
+            self.assertIs(df_shuffled_2["str"].tolist() == filtered_str, bool_to_check, msg=msg)
+            self.assertIs(series_num_shuffled_2.tolist() == filtered_num, bool_to_check, msg=msg)
+
 
 if __name__ == "__main__":
     ut.main()
diff --git a/tests/test_metrics.py b/tests/test_metrics.py
index bdffd6409..af728c14e 100644
--- a/tests/test_metrics.py
+++ b/tests/test_metrics.py
@@ -10,8 +10,8 @@
     Thresholder,
     pearsonr,
     spearmanr,
-    mean_squared_error,
 )
+from torchmetrics.functional import mean_squared_error
 
 
 class test_Metrics(ut.TestCase):
diff --git a/tests/test_predictor.py b/tests/test_predictor.py
index 275526c78..192a98b12 100644
--- a/tests/test_predictor.py
+++ b/tests/test_predictor.py
@@ -16,7 +16,6 @@ def test_parse_loss_fun(self):
         target = (torch.rand(10, 5) > 0.5).to(preds.dtype)
         for this_loss in losses:
             loss_fun = PredictorModule.parse_loss_fun(this_loss)
-            print(this_loss)
             loss = loss_fun(preds, target)
 
 
diff --git a/tests/test_predictor_module.py b/tests/test_predictor_module.py
deleted file mode 100644
index deb463bd2..000000000
--- a/tests/test_predictor_module.py
+++ /dev/null
@@ -1,50 +0,0 @@
-import yaml
-
-import goli
-
-
-# def test_load_pretrained_model():
-#     predictor = goli.trainer.PredictorModule.load_pretrained_models("goli-zinc-micro-dummy-test")
-#     assert isinstance(predictor, goli.trainer.predictor.PredictorModule)
-
-
-# def test_training(datadir, tmpdir):
-
-#     config_path = datadir / "config_micro_ZINC.yaml"
-#     data_path = datadir / "micro_ZINC.csv"
-
-#     # Load a config
-#     with open(config_path, "r") as file:
-#         yaml_config = yaml.load(file, Loader=yaml.FullLoader)
-
-#     training_dir = tmpdir.mkdir("training")
-
-#     # Tweak config and paths
-#     yaml_config["datamodule"]["args"]["df_path"] = data_path
-#     yaml_config["datamodule"]["args"]["cache_data_path"] = None
-
-#     yaml_config["trainer"]["trainer"]["min_epochs"] = 1
-#     yaml_config["trainer"]["trainer"]["max_epochs"] = 1
-
-#     yaml_config["trainer"]["logger"]["save_dir"] = training_dir
-#     yaml_config["trainer"]["model_checkpoint"]["dirpath"] = None
-#     yaml_config["trainer"]["trainer"]["default_root_dir"] = training_dir
-#     yaml_config["trainer"]["trainer"]["gpus"] = 0
-#     yaml_config["trainer"]["trainer"]["limit_train_batches"] = 1
-#     yaml_config["trainer"]["trainer"]["limit_val_batches"] = 1
-#     yaml_config["trainer"]["trainer"]["limit_test_batches"] = 1
-
-#     # Load datamodule
-#     datamodule = goli.config.load_datamodule(yaml_config)
-
-#     # Load a trainer
-#     trainer = goli.config.load_trainer(yaml_config)
-
-#     # Load a pretrained model
-#     predictor = goli.trainer.PredictorModule.load_pretrained_models("goli-zinc-micro-dummy-test")
-
-#     # Inference
-#     results = trainer.predict(predictor, datamodule=datamodule, return_predictions=True)
-
-#     assert len(results) == 8  # type: ignore
-#     assert tuple(results[0].shape) == (128, 1)  # type: ignore
diff --git a/tests/test_utils.py b/tests/test_utils.py
index 4cdd8427b..4b3a99c5f 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -2,7 +2,6 @@
 Unit tests for the metrics and wrappers of goli/utils/...
 """
 
-from typing import Tuple
 from goli.utils.tensor import nan_mad, nan_mean, nan_std, nan_var, nan_median
 import torch
 import numpy as np