code: def setup(self, stage: Optional[str] = None):
      t = time.perf_counter()
      print("Reading dataset...", end=" ", flush=True)
      dataset = MAG240MDataset(self.data_dir)
  
      self.train_idx = torch.from_numpy(dataset.get_idx_split("train"))
      self.train_idx = self.train_idx
      self.train_idx.share_memory_()
      self.val_idx = torch.from_numpy(dataset.get_idx_split("valid"))
      self.val_idx.share_memory_()
      self.test_idx = torch.from_numpy(
          dataset.get_idx_split("test-dev")
      )
      self.test_idx.share_memory_()
  
      N = (
          dataset.num_papers
          + dataset.num_authors
          + dataset.num_institutions
      )
  
      x = np.memmap(
          f"{dataset.dir}/full_feat.npy",
          dtype=np.float16,
          mode="r",
          shape=(N, self.num_features),
      )
  
      if self.in_memory:
          self.x = np.empty((N, self.num_features), dtype=np.float16)
          self.x[:] = x
          self.x = torch.from_numpy(self.x).share_memory_()
      else:
          self.x = x
  
      self.y = torch.from_numpy(dataset.all_paper_label)
  
      path = f"{dataset.dir}/full_adj_t.pt"
      self.adj_t = torch.load(path)
      print(f"Done! [{time.perf_counter() - t:.2f}s]")
  
imports: import argparse
  import glob
  import os
  import os.path
  import time
  from typing import List, NamedTuple, Optional
  import numpy
  import torch
  import torch.nn.functional
  from ogb.lsc import MAG240MDataset, MAG240MEvaluator
  import pytorch_lightning
  from pytorch_lightning import LightningDataModule, LightningModule, Trainer, seed_everything
  from pytorch_lightning.callbacks import ModelCheckpoint
  from torchmetrics import Accuracy
  from torch import Tensor
  from torch.nn import BatchNorm1d, Dropout, Linear, ModuleList, ReLU, Sequential
  from torch.optim.lr_scheduler import StepLR
  from torch_geometric.data import NeighborSampler
  from torch_geometric.nn import GATConv, SAGEConv
  from torch_sparse import SparseTensor
  from tqdm.auto import tqdm
  from root import ROOT
label_desc: Use image node_id 6 for calling the MAG240M obj's underlying member method code with example usage: obj.setup(stage) without return types