code: def gen_masks(
      y: Tensor,
      train_per_class: int = 20,
      val_per_class: int = 30,
      num_splits: int = 20,
  ) -> Tuple[Tensor, Tensor, Tensor]:
      num_classes = int(y.max()) + 1
  
      train_mask = torch.zeros(y.size(0), num_splits, dtype=torch.bool)
      val_mask = torch.zeros(y.size(0), num_splits, dtype=torch.bool)
  
      for c in range(num_classes):
          idx = (y == c).nonzero(as_tuple=False).view(-1)
          perm = torch.stack(
              [torch.randperm(idx.size(0)) for _ in range(num_splits)],
              dim=1,
          )
          idx = idx[perm]
  
          train_idx = idx[:train_per_class]
          train_mask.scatter_(0, train_idx, True)
          val_idx = idx[
              train_per_class : train_per_class + val_per_class
          ]
          val_mask.scatter_(0, val_idx, True)
  
      test_mask = ~(train_mask | val_mask)
  
      return train_mask, val_mask, test_mask
  
file_path: DIG/dig/lsgraph/method/GraphFMOB/utils.py
imports: from typing import Optional, Tuple
  import torch
  from torch import Tensor
  import torch.nn.functional
  from torch_sparse import SparseTensor
label_desc: Use image node_id 2 for calling a global function with example usage: gen_masks(y, train_per_class, val_per_class, num_splits) and returns: train_mask,  val_mask,  test_mask