code: def training_step(self, batch, batch_idx):
      blocks = [block.to("cuda") for block in batch.blocks]
      x = batch.node_features["feat"]
      y = batch.labels.to("cuda")
      y_hat = self(blocks, x)
      loss = F.cross_entropy(y_hat, y)
      self.train_acc(torch.argmax(y_hat, 1), y)
      self.log(
          "train_acc",
          self.train_acc,
          prog_bar=True,
          on_step=True,
          on_epoch=False,
      )
      self.log_node_and_edge_counts(blocks)
      return loss
  
imports: import argparse
  import dgl.graphbolt
  import dgl.nn.pytorch
  import torch
  import torch.nn
  import torch.nn.functional
  from pytorch_lightning import LightningDataModule, LightningModule, Trainer
  from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
  from torchmetrics import Accuracy
label_desc: Use image node_id 4 for calling the SAGE obj's underlying member method code with example usage: obj.training_step(batch, batch_idx) and returns: loss