code: def forward(self, data_batch):
      z, pos, batch = (
          data_batch["atom_type"],
          data_batch["position"],
          data_batch["batch"],
      )
      node_feat = self.feat_net(z, pos, batch)
      focus_score = self.focus_mlp(node_feat)
  
      new_atom_type, focus = (
          data_batch["new_atom_type"],
          data_batch["focus"],
      )
      x_z = F.one_hot(
          new_atom_type, num_classes=self.num_node_types
      ).float()
      x_z += self.deq_coeff * torch.rand(x_z.size(), device=x_z.device)
  
      local_node_type_feat, query_batch = (
          node_feat[focus[:, 0]],
          batch[focus[:, 0]],
      )
      global_node_type_feat = self.node_att(
          local_node_type_feat, node_feat, node_feat, query_batch, batch
      )
      node_type_feat = torch.cat(
          (local_node_type_feat, global_node_type_feat), dim=-1
      )
      node_latent, node_log_jacob = flow_forward(
          self.node_flow_layers, x_z, node_type_feat
      )
  
      node_type_emb_block = self.feat_net.init_e.emb
      node_type_emb = node_type_emb_block(new_atom_type)
      node_emb = node_feat * node_type_emb[batch]
  
      c1_focus, c2_c1_focus = (
          data_batch["c1_focus"],
          data_batch["c2_c1_focus"],
      )
      dist, angle, torsion = (
          data_batch["new_dist"],
          data_batch["new_angle"],
          data_batch["new_torsion"],
      )
  
      local_dist_feat, dist_query_batch = (
          node_emb[focus[:, 0]],
          batch[focus[:, 0]],
      )
      global_dist_feat = self.dist_att(
          local_dist_feat, node_emb, node_emb, dist_query_batch, batch
      )
      local_angle_feat, angle_query_batch = (
          torch.cat(
              (node_emb[c1_focus[:, 1]], node_emb[c1_focus[:, 0]]),
              dim=1,
          ),
          batch[c1_focus[:, 0]],
      )
      global_angle_feat = self.angle_att(
          local_angle_feat, node_emb, node_emb, angle_query_batch, batch
      )
      local_torsion_feat, torsion_query_batch = (
          torch.cat(
              (
                  node_emb[c2_c1_focus[:, 2]],
                  node_emb[c2_c1_focus[:, 1]],
                  node_emb[c2_c1_focus[:, 0]],
              ),
              dim=1,
          ),
          batch[c2_c1_focus[:, 0]],
      )
      global_torsion_feat = self.torsion_att(
          local_torsion_feat,
          node_emb,
          node_emb,
          torsion_query_batch,
          batch,
      )
      dist_feat = torch.cat((local_dist_feat, global_dist_feat), dim=-1)
      angle_feat = torch.cat(
          (local_angle_feat, global_angle_feat), dim=-1
      )
      torsion_feat = torch.cat(
          (local_torsion_feat, global_torsion_feat), dim=-1
      )
  
      dist_latent, dist_log_jacob = flow_forward(
          self.dist_flow_layers, dist, dist_feat
      )
      angle_latent, angle_log_jacob = flow_forward(
          self.angle_flow_layers, angle, angle_feat
      )
      torsion_latent, torsion_log_jacob = flow_forward(
          self.torsion_flow_layers, torsion, torsion_feat
      )
  
      return (
          (node_latent, node_log_jacob),
          focus_score,
          (dist_latent, dist_log_jacob),
          (angle_latent, angle_log_jacob),
          (torsion_latent, torsion_log_jacob),
      )
  
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 2 for calling the SphGen obj's underlying member method code with example usage: obj.forward(data_batch) and returns: node_latent,  node_log_jacob,  focus_score,  dist_latent,  dist_log_jacob,  angle_latent,  angle_log_jacob,  torsion_latent,  torsion_log_jacob