code: def __init__(self, device_idx):
      super().__init__()
      self.handle = pynvml.nvmlDeviceGetHandleByIndex(device_idx)
  
file_path: flash-attention/training/src/utils/gpu_affinity.py
imports: import collections
  import math
  import os
  import pathlib
  import re
  import pynvml
label_desc: Use image node_id 1 to create a new device object  with example: obj = device(device_idx)