code: def compute_loss(
      cls_logits: Tensor,
      labels: Tensor,
      logits_s1: Tensor,
      logits_s2: Tensor,
      epoch: int,
      total_epochs: int,
      device: torch.device,
  ):
      # classification loss
      classify_loss = F.nll_loss(cls_logits, labels.to(device))
  
      # loss for vertex infomax pooling
      scale1, scale2 = logits_s1.size(0) // 2, logits_s2.size(0) // 2
      s1_label_t, s1_label_f = torch.ones(scale1), torch.zeros(scale1)
      s2_label_t, s2_label_f = torch.ones(scale2), torch.zeros(scale2)
      s1_label = torch.cat((s1_label_t, s1_label_f), dim=0).to(device)
      s2_label = torch.cat((s2_label_t, s2_label_f), dim=0).to(device)
  
      pool_loss_s1 = F.binary_cross_entropy_with_logits(
          logits_s1, s1_label
      )
      pool_loss_s2 = F.binary_cross_entropy_with_logits(
          logits_s2, s2_label
      )
      pool_loss = (pool_loss_s1 + pool_loss_s2) / 2
  
      loss = classify_loss + (2 - epoch / total_epochs) * pool_loss
  
      return loss
  
file_path: dgl/examples/pytorch/gxn/main_early_stop.py
imports: import json
  import os
  from datetime import datetime
  from time import time
  import dgl
  import torch
  import torch.nn.functional
  from data_preprocess import degree_as_feature, node_label_as_feature
  from dgl.data import LegacyTUDataset
  from dgl.dataloading import GraphDataLoader
  from networks import GraphClassifier
  from torch import Tensor
  from torch.utils.data import random_split
  from utils import get_stats, parse_args
label_desc: Use image node_id 1 for calling a global function with example usage: compute_loss(cls_logits, labels, logits_s1, logits_s2, epoch, total_epochs, device) and returns: loss