cuda_test.py

import torch


if __name__ == "__main__":
    print(torch.cuda.is_available())

    device_count = torch.cuda.device_count()

    if device_count > 0:
        for i in range(device_count):
            device_name = torch.cuda.get_device_name(i)
            print(f"cuda:{i}: {device_name}")
    else:
        print("No CUDA devices available.")