code: def initalize_pytorch_worker(dev_id):
      import cupy as cp
      import rmm
      from rmm.allocators.torch import rmm_torch_allocator
      from rmm.allocators.cupy import rmm_cupy_allocator
  
      dev = cp.cuda.Device(
          dev_id
      )  # Create cuda context on the right gpu, defaults to gpu-0
      dev.use()
      rmm.reinitialize(
          pool_allocator=True,
          initial_pool_size=10e9,
          maximum_pool_size=15e9,
          devices=[dev_id],
      )
  
      if dev_id == 0:
          torch.cuda.memory.change_current_allocator(
              rmm_torch_allocator
          )
  
      torch.cuda.set_device(dev_id)
      cp.cuda.set_allocator(rmm_cupy_allocator)
      enable_spilling()
      print("device_id", dev_id, flush=True)
  
file_path: cugraph/python/cugraph-dgl/examples/multi_trainer_MG_example/workflow.py
imports: import dgl
  import torch
  import time
  from distributed import Client, Event
  import tempfile
  from cugraph.dask.comms import comms
label_desc: Use image node_id 4 for calling a global function with example usage: initalize_pytorch_worker(dev_id) without return types