code: def __init__(self, in_channel=256, num_unit=8, device="cpu"):
      super(PrimaryCapsuleLayer, self).__init__()
      self.in_channel = in_channel
      self.num_unit = num_unit
      self.deivce = device
      self.conv_units = nn.ModuleList(
          [
              nn.Conv2d(self.in_channel, 32, 9, 2)
              for _ in range(self.num_unit)
          ]
      )
  
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 1 to create a new PrimaryCapsuleLayer object from inherited base classes: nn with example: obj = PrimaryCapsuleLayer(in_channel, num_unit, device)