import torch
import torch.nn as nn

# 查询设备和显存信息
def query_device():
    if torch.cuda.is_available():
        print(f"CUDA is available. Using {torch.cuda.device_count()} device(s).")
        for i in range(torch.cuda.device_count()):
            print(f"Device {i}: {torch.cuda.get_device_name(i)}")
            print(f"  Memory Allocated: {torch.cuda.memory_allocated(i) / (1024 ** 2):.2f} MB")
            print(f"  Memory Cached: {torch.cuda.memory_reserved(i) / (1024 ** 2):.2f} MB")
    else:
        print("CUDA is not available. Using CPU.")

# 模拟显存分配，直到触发 OOM
def trigger_oom(device):
    try:
        print("\nStarting memory allocation...")
        tensor_list = []
        while True:
            # 每次分配一个大的张量来消耗显存
            tensor = torch.randn((10000, 10000), device=device)
            tensor_list.append(tensor)
            # 打印当前显存状态
            allocated_memory = torch.cuda.memory_allocated(device) / (1024 ** 2)  # MB
            cached_memory = torch.cuda.memory_reserved(device) / (1024 ** 2)  # MB
            print(f"Allocated Memory: {allocated_memory:.2f} MB, Cached Memory: {cached_memory:.2f} MB")
    except RuntimeError as e:
        if "out of memory" in str(e):
            print(f"OOM Triggered! {e}")
        else:
            print(f"Unexpected error: {e}")

# 主程序
if __name__ == "__main__":
    # 查询设备信息
    query_device()

    # 设置设备为 CUDA（如果可用）
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    if torch.cuda.is_available():
        # 在设备上触发 OOM
        trigger_oom(device)
    else:
        print("CUDA not available. Unable to trigger OOM.")