code: def k_fold(dataset, folds):
      skf = StratifiedKFold(folds, shuffle=True, random_state=12345)
  
      test_indices, train_indices = [], []
      for _, idx in skf.split(
          torch.zeros(len(dataset)), dataset.data.y
      ):
          test_indices.append(torch.from_numpy(idx).to(torch.long))
  
      val_indices = [test_indices[i - 1] for i in range(folds)]
  
      for i in range(folds):
          train_mask = torch.ones(len(dataset), dtype=torch.bool)
          train_mask[test_indices[i]] = 0
          train_mask[val_indices[i]] = 0
          train_indices.append(
              train_mask.nonzero(as_tuple=False).view(-1)
          )
  
      return train_indices, test_indices, val_indices
  
file_path: pytorch_geometric/benchmark/kernel/train_eval.py
imports: import time
  import torch
  import torch.nn.functional
  from sklearn.model_selection import StratifiedKFold
  from torch import tensor
  from torch.optim import Adam
  from torch_geometric.loader import DataLoader
  from torch_geometric.loader import DenseDataLoader
label_desc: Use image node_id 2 for calling a global function with example usage: k_fold(dataset, folds) and returns: train_indices,  test_indices,  val_indices