code: def __init__(
      self,
      cutoff,
      num_node_types,
      num_layers,
      hidden_channels,
      int_emb_size,
      basis_emb_size,
      out_emb_channels,
      num_spherical,
      num_radial,
      num_flow_layers,
      deq_coeff=0.9,
      use_gpu=True,
      n_att_heads=4,
  ):
      super(SphGen, self).__init__()
      self.use_gpu = use_gpu
      self.num_node_types = num_node_types
  
      self.feat_net = SphereNet(
          cutoff,
          num_node_types,
          num_layers,
          hidden_channels,
          int_emb_size,
          basis_emb_size,
          out_emb_channels,
          num_spherical,
          num_radial,
      )
      node_feat_dim, dist_feat_dim, angle_feat_dim, torsion_feat_dim = (
          hidden_channels * 2,
          hidden_channels * 2,
          hidden_channels * 3,
          hidden_channels * 4,
      )
  
      self.node_flow_layers = nn.ModuleList(
          [
              ST_Net_Exp(
                  node_feat_dim,
                  num_node_types,
                  hid_dim=hidden_channels,
                  bias=True,
              )
              for _ in range(num_flow_layers)
          ]
      )
      self.dist_flow_layers = nn.ModuleList(
          [
              ST_Net_Exp(
                  dist_feat_dim, 1, hid_dim=hidden_channels, bias=True
              )
              for _ in range(num_flow_layers)
          ]
      )
      self.angle_flow_layers = nn.ModuleList(
          [
              ST_Net_Exp(
                  angle_feat_dim, 1, hid_dim=hidden_channels, bias=True
              )
              for _ in range(num_flow_layers)
          ]
      )
      self.torsion_flow_layers = nn.ModuleList(
          [
              ST_Net_Exp(
                  torsion_feat_dim,
                  1,
                  hid_dim=hidden_channels,
                  bias=True,
              )
              for _ in range(num_flow_layers)
          ]
      )
      self.focus_mlp = MLP(hidden_channels)
      self.deq_coeff = deq_coeff
  
      self.node_att = MH_ATT(
          n_att_heads,
          q_dim=hidden_channels,
          k_dim=hidden_channels,
          v_dim=hidden_channels,
          out_dim=hidden_channels,
      )
      self.dist_att = MH_ATT(
          n_att_heads,
          q_dim=hidden_channels,
          k_dim=hidden_channels,
          v_dim=hidden_channels,
          out_dim=hidden_channels,
      )
      self.angle_att = MH_ATT(
          n_att_heads,
          q_dim=2 * hidden_channels,
          k_dim=hidden_channels,
          v_dim=hidden_channels,
          out_dim=hidden_channels,
      )
      self.torsion_att = MH_ATT(
          n_att_heads,
          q_dim=3 * hidden_channels,
          k_dim=hidden_channels,
          v_dim=hidden_channels,
          out_dim=hidden_channels,
      )
      if use_gpu:
          (
              self.node_att,
              self.dist_att,
              self.angle_att,
              self.torsion_att,
          ) = (
              self.node_att.to("cuda"),
              self.dist_att.to("cuda"),
              self.angle_att.to("cuda"),
              self.torsion_att.to("cuda"),
          )
  
      if use_gpu:
          self.feat_net = self.feat_net.to("cuda")
          self.node_flow_layers = self.node_flow_layers.to("cuda")
          self.dist_flow_layers = self.dist_flow_layers.to("cuda")
          self.angle_flow_layers = self.angle_flow_layers.to("cuda")
          self.torsion_flow_layers = self.torsion_flow_layers.to("cuda")
          self.focus_mlp = self.focus_mlp.to("cuda")
  
file_path: DIG/dig/ggraph3D/method/G_SphereNet/model/sphgen.py
imports: import torch
  import torch.nn
  import torch.nn.functional
  from .spherenet import SphereNet
  from .net_utils import *
  from .geometric_computing import *
  from .att import MH_ATT
label_desc: Use image node_id 1 to create a new SphGen object from inherited base classes: nn with example: obj = SphGen(cutoff, num_node_types, num_layers, hidden_channels, int_emb_size, basis_emb_size, out_emb_channels, num_spherical, num_radial, num_flow_layers, deq_coeff, use_gpu, n_att_heads)