def METHOD_NAME(self, g, inputs):
    h = inputs
    for i, layer in enumerate(self.gat_layers):
        h = layer(g, h)
        if i == 1:  # last layer
            h = h.mean(1)
        else:  # other layer(s)
            h = h.flatten(1)
    return h