Update gin.py

Author: cthoyt

torchdrug/models/gin.py

@@ -26,7 +26,7 @@ class GraphIsomorphismNetwork(nn.Module, core.Configurable):
         batch_norm (bool, optional): apply batch normalization or not
         activation (str or function, optional): activation function
         concat_hidden (bool, optional): concat hidden representations from all layers as output
-        readout (str, optional): readout function. Available functions are ``sum`` and ``mean``.
+        readout (str, optional): readout function. Available functions are ``sum``, ``mean``, and ``max``.
     """
 
     def __init__(self, input_dim=None, hidden_dims=None, edge_input_dim=None, num_mlp_layer=2, eps=0, learn_eps=False,
@@ -52,6 +52,8 @@ def __init__(self, input_dim=None, hidden_dims=None, edge_input_dim=None, num_ml
             self.readout = layers.SumReadout()
         elif readout == "mean":
             self.readout = layers.MeanReadout()
+        elif readout == "max":
+            self.readout = layers.MaxReadout()
         else:
             raise ValueError("Unknown readout `%s`" % readout)
 
@@ -88,4 +90,4 @@ def forward(self, graph, input, all_loss=None, metric=None):
         return {
             "graph_feature": graph_feature,
             "node_feature": node_feature
-        }
+        }