code: def HAVE_AMPERE_GPU():
      if HAVE_PYNVML:
          pynvml.nvmlInit()
          handle = pynvml.nvmlDeviceGetHandleByIndex(0)
          device_arch = pynvml.nvmlDeviceGetArchitecture(handle)
          pynvml.nvmlShutdown()
          return device_arch == pynvml.NVML_DEVICE_ARCH_AMPERE
      else:
          return False
  
file_path: NeMo/tests/collections/nlp/test_flash_attention.py
imports: import random
  import pytest
  import torch
  from megatron.core import ModelParallelConfig
  from pytorch_lightning.trainer.trainer import Trainer
  from nemo.collections.nlp.modules.common.megatron.attention import CoreAttention
  from nemo.collections.nlp.modules.common.megatron.megatron_init import initialize_model_parallel_for_nemo
  from nemo.collections.nlp.modules.common.megatron.utils import build_attention_mask_3d
  from nemo.collections.nlp.parts.nlp_overrides import NLPDDPStrategy
label_desc: Use image node_id 1 for calling a global function with example usage: HAVE_AMPERE_GPU() and returns: unknown, False