code: def forward(self, x, adj, log=False):
      z_l = self.embed(x, adj)
      s_l = self.assign(x, adj)
      if log:
          self.log["s"] = s_l.cpu().numpy()
      xnext = torch.matmul(s_l.transpose(-1, -2), z_l)
      anext = (s_l.transpose(-1, -2)).matmul(adj).matmul(s_l)
  
      for loss_layer in self.reg_loss:
          loss_name = str(type(loss_layer).__name__)
          self.loss_log[loss_name] = loss_layer(adj, anext, s_l)
      if log:
          self.log["a"] = anext.cpu().numpy()
      return xnext, anext
  
file_path: dgl/examples/pytorch/diffpool/model/tensorized_layers/diffpool.py
imports: import torch
  from torch import nn
  from model.loss import EntropyLoss, LinkPredLoss
  from model.tensorized_layers.assignment import DiffPoolAssignment
  from model.tensorized_layers.graphsage import BatchedGraphSAGE
label_desc: Use image node_id 2 for calling the BatchedDiffPool obj's underlying member method code with example usage: obj.forward(x, adj, log) and returns: xnext,  anext