code: def __init__(
      self,
      in_nodes_dim=8,
      in_nodes=1152,
      out_nodes=10,
      out_nodes_dim=16,
      device="cpu",
  ):
      super(DGLDigitCapsuleLayer, self).__init__()
      self.device = device
      self.in_nodes_dim, self.out_nodes_dim = (
          in_nodes_dim,
          out_nodes_dim,
      )
      self.in_nodes, self.out_nodes = in_nodes, out_nodes
      self.weight = nn.Parameter(
          torch.randn(in_nodes, out_nodes, out_nodes_dim, in_nodes_dim)
      )
  
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 1 to create a new DGLDigitCapsuleLayer object from inherited base classes: nn with example: obj = DGLDigitCapsuleLayer(in_nodes_dim, in_nodes, out_nodes, out_nodes_dim, device)