class Net(torch.nn.Module):
    def __init__(self, num_relations, num_classes, num_nodes=None, input_dim=None, hidden_dim=16, num_bases=30):
        super().__init__()
        assert num_nodes is not None or input_dim is not None, "Please provide input feature dimensionality or number of nodes"
        self.conv1 = RGCNConv(num_nodes if input_dim is None else input_dim, hidden_dim, num_relations,
                              num_bases)
        self.conv2 = RGCNConv(hidden_dim, num_classes, dataset.num_relations,
                              num_bases)

    def forward(self, x, edge_index, edge_type):
        # if x is None, uses an embedding based on num_nodes
        x = F.relu(self.conv1(x, edge_index, edge_type))
        x = self.conv2(x, edge_index, edge_type)
        return F.log_softmax(x, dim=1)