revise import issue for torch-scatter in grin and bfgnn

Author: liugangcode

README.md

@@ -56,6 +56,8 @@
 | Model | Required Packages |
 |-------|-------------------|
 | HFPretrainedMolecularEncoder | transformers |
+| BFGNNMolecularPredictor | torch-scatter |
+| GRINMolecularPredictor | torch-scatter |
 
 ## Usage
 

torch_molecule/generator/pretrained/modeling_pretrained.py

@@ -1,3 +1 @@
-# TODO
-
-# safe-100m: https://huggingface.co/anrilombard/safe-100m
+# TODO

torch_molecule/nn/gnn.py

@@ -4,7 +4,12 @@
 from torch_geometric.nn.norm import GraphNorm, PairNorm, DiffGroupNorm, InstanceNorm, LayerNorm, GraphSizeNorm
 from torch_geometric.nn import MessagePassing
 from torch_geometric.nn import global_add_pool
-from torch_scatter import scatter_min, scatter_max, scatter_mean,scatter_add
+try:
+    from torch_scatter import scatter_min, scatter_max, scatter_mean, scatter_add
+    _has_torch_scatter = True
+except ImportError:
+    scatter_min = scatter_max = scatter_mean = scatter_add = None
+    _has_torch_scatter = False
 
 from ..utils import get_atom_feature_dims, get_bond_feature_dims 
 
@@ -192,6 +197,9 @@ def message(self, x_j, edge_attr):
         return self.f_agg(torch.cat([x_j, edge_attr], dim=1))
 
     def aggregate(self, inputs, index, dim_size=None):
+        if not _has_torch_scatter or scatter_min is None:
+            raise ImportError("BFGNN requires `torch_scatter` package. Please install it via `pip install torch-scatter -f https://data.pyg.org/whl/torch-${TORCH}+${CUDA}.html`.")
+
         out, _ = scatter_min(inputs, index, dim=0, dim_size=dim_size)
         out[out == float('inf')] = 0.0
         return out
@@ -240,6 +248,9 @@ def message(self, x_j, edge_attr, norm):
         return norm.view(-1, 1) * F.relu(m)
 
     def aggregate(self, inputs, index, dim_size=None):
+        if not _has_torch_scatter or scatter_min is None:
+            raise ImportError("BFGNN requires `torch_scatter` package. Please install it via `pip install torch-scatter -f https://data.pyg.org/whl/torch-${TORCH}+${CUDA}.html`.")
+
         out, _ = scatter_min(inputs, index, dim=0, dim_size=dim_size)
         out[out == float('inf')] = 0.0
         return out
@@ -285,6 +296,9 @@ def message(self, x_j, edge_attr):
         return F.relu(x_j + edge_attr)
 
     def aggregate(self, inputs, index, dim_size=None):
+        if not _has_torch_scatter or scatter_max is None:
+            raise ImportError("GRIN requires `torch_scatter` package. Please install it via `pip install torch-scatter -f https://data.pyg.org/whl/torch-${TORCH}+${CUDA}.html`.")
+
         out, _ = scatter_max(inputs, index, dim=0, dim_size=dim_size)
         out[out == float('inf')] = 0.0
         return out
@@ -339,6 +353,9 @@ def message(self, x_j, edge_attr, norm):
         return norm.view(-1,1) * F.relu(x_j + edge_attr)
 
     def aggregate(self, inputs, index, dim_size=None):
+        if not _has_torch_scatter or scatter_max is None:
+            raise ImportError("GRIN requires `torch_scatter` package. Please install it via `pip install torch-scatter -f https://data.pyg.org/whl/torch-${TORCH}+${CUDA}.html`.")
+
         out, _ = scatter_max(inputs, index, dim=0, dim_size=dim_size)
         out[out == float('inf')] = 0.0
         return out