import torch
import time

def keep_all_gpus_busy():
    if not torch.cuda.is_available():
        print("No CUDA devices available.")
        return

    num_gpus = torch.cuda.device_count()
    print(f"[keepalive] Using {num_gpus} GPUs")

    tensors = []
    for gpu_id in range(num_gpus):
        device = torch.device(f"cuda:{gpu_id}")
        x = torch.randn(60120, 60120, device=device)
        tensors.append((device, x))

    while True:
        for device, x in tensors:
            with torch.cuda.device(device):
                x = torch.matmul(x, x)
                torch.cuda.synchronize(device)
        time.sleep(2)

if __name__ == "__main__":
    import torch

    print("torch version:", torch.__version__)
    print("HIP version:", torch.version.hip)
    print("CUDA available:", torch.cuda.is_available())

    if torch.cuda.is_available():
        print(f"Device count: {torch.cuda.device_count()}")
        for i in range(torch.cuda.device_count()):
            print(f"Device {i}: {torch.cuda.get_device_name(i)}")
    else:
        print("No GPU available (NVIDIA or AMD ROCm).")
        
    keep_all_gpus_busy()