import torch
print(torch.cuda.is_available())

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "2"