cuda_test.py

def main():
    import torch
    print(torch.cuda.is_available())
    print(torch.cuda.get_device_name())

if __name__ == '__main__':
    main()