code: def init_pytorch_worker(
      rank: int, use_rmm_torch_allocator: bool = False
  ) -> None:
      import cupy
      import rmm
      from pynvml.smi import nvidia_smi
  
      smi = nvidia_smi.getInstance()
      pool_size = 16e9  # FIXME calculate this
  
      rmm.reinitialize(
          devices=[rank],
          pool_allocator=True,
          initial_pool_size=pool_size,
      )
  
      if use_rmm_torch_allocator:
          warnings.warn(
              "Using the rmm pytorch allocator is currently unsupported."
              " The default allocator will be used instead."
          )
          # FIXME somehow get the pytorch allocator to work
          # from rmm.allocators.torch import rmm_torch_allocator
          # torch.cuda.memory.change_current_allocator(rmm_torch_allocator)
  
      from rmm.allocators.cupy import rmm_cupy_allocator
  
      cupy.cuda.set_allocator(rmm_cupy_allocator)
  
      cupy.cuda.Device(rank).use()
      torch.cuda.set_device(rank)
  
      # Pytorch training worker initialization
      torch.distributed.init_process_group(backend="nccl")
  
file_path: cugraph/benchmarks/cugraph/standalone/bulk_sampling/bench_cugraph_training.py
imports: import os
  import argparse
  import json
  import warnings
  import torch
  import numpy
  import pandas
  import torch.distributed
  from datasets import OGBNPapers100MDataset
  from cugraph.testing.mg_utils import enable_spilling
label_desc: Use image node_id 1 for calling a global function with example usage: init_pytorch_worker(rank, use_rmm_torch_allocator) without return types