code: def init_pytorch_worker(device_id: int) -> None:
      import cupy
      import rmm
  
      rmm.reinitialize(
          devices=[device_id],
          pool_allocator=False,
      )
  
      cupy.cuda.Device(device_id).use()
      torch.cuda.set_device(device_id)
  
file_path: cugraph/python/cugraph-pyg/cugraph_pyg/examples/graph_sage_sg.py
imports: import time
  import argparse
  import gc
  import torch
  from torch_geometric.nn import CuGraphSAGEConv
  import torch.nn
  import torch.nn.functional
  from typing import Union
label_desc: Use image node_id 1 for calling a global function with example usage: init_pytorch_worker(device_id) without return types