categories_example.py

# here's a dynamic reduction network that can categorize
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.drn = DynamicReductionNetwork(input_dim=3, hidden_dim=64,
                                           k = 16,
                                           output_dim=2, aggr='add',
                                           norm=torch.tensor([1., 1./27., 1./27.]))
        
    def forward(self, data):
        logits = self.drn(data)
        return F.log_softmax(logits, dim=1)

# here's the training setup
model = Net().to(device)
for datum in data:
  datum = datum.to(device)
  result = model(datum)
  F.nll_loss(result, datum).backward()
  optimizer.step()