test_df.py

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()