code: def get_nvml_gpu_id(torch_gpu_id):
      """
      Remap torch device id to nvml device id, respecting CUDA_VISIBLE_DEVICES.
  
      If the latter isn't set return the same id
      """
      # if CUDA_VISIBLE_DEVICES is used automagically remap the id since pynvml ignores this env var
      if "CUDA_VISIBLE_DEVICES" in os.environ:
          ids = list(
              map(
                  int,
                  os.environ.get("CUDA_VISIBLE_DEVICES", "").split(","),
              )
          )
          return ids[torch_gpu_id]  # remap
      else:
          return torch_gpu_id
  
imports: from IPython import get_ipython
  from collections import namedtuple
  import datetime
  import gc
  import logging
  import os
  import psutil
  import random
  import sys
  import threading
  import time
label_desc: Use image node_id 4 for calling a global function with example usage: get_nvml_gpu_id(torch_gpu_id) and returns: ids, torch_gpu_id