code: def forward(self, x):
      self.batch_size = x.size(0)
      u_hat = self.compute_uhat(x)
      routing = DGLRoutingLayer(
          self.in_nodes,
          self.out_nodes,
          self.out_nodes_dim,
          batch_size=self.batch_size,
          device=self.device,
      )
      routing(u_hat, routing_num=3)
      out_nodes_feature = routing.g.nodes[routing.out_indx].data["v"]
      # shape transformation is for further classification
      return (
          out_nodes_feature.transpose(0, 1)
          .unsqueeze(1)
          .unsqueeze(4)
          .squeeze(1)
      )
  
file_path: dgl/examples/pytorch/capsule/DGLDigitCapsule.py
imports: import dgl
  import dgl.function
  import torch
  from DGLRoutingLayer import DGLRoutingLayer
  from torch import nn
  from torch.nn import functional
label_desc: Use image node_id 2 for calling the DGLDigitCapsuleLayer obj's underlying member method code with example usage: obj.forward(x) and returns: out_nodes_feature