import torch def check_bf16_support(): if not torch.cuda.is_available(): print("CUDA is not available on this system.") return False device = torch.device("cuda") capability = torch.cuda.get_device_capability(device) # As of now, GPUs with compute capability >= 8.0 support BF16 # Example: NVIDIA A100 has compute capability 8.0 bf16_supported = capability[0] >= 8 print(f"GPU Compute Capability: {capability}") if bf16_supported: print("BF16 is supported on this GPU.") else: print("BF16 is not supported on this GPU.") return bf16_supported # Check if BF16 is supported check_bf16_support()