code: def deviceGetCount():
      return pynvml.nvmlDeviceGetCount()
  
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 2 for calling a global function with example usage: deviceGetCount() and returns: pynvml