code: def forward(self, x):
      unit = [
          self.conv_units[i](x) for i, l in enumerate(self.conv_units)
      ]
      unit = torch.stack(unit, dim=1)
      batch_size = x.size(0)
      unit = unit.view(batch_size, 8, -1)
      return squash(unit, dim=2)
  
file_path: dgl/examples/pytorch/capsule/model.py
imports: import torch
  from DGLDigitCapsule import DGLDigitCapsuleLayer
  from DGLRoutingLayer import squash
  from torch import nn
label_desc: Use image node_id 2 for calling the PrimaryCapsuleLayer obj's underlying member method code with example usage: obj.forward(x) and returns: squash