就用cpu训练网络

This commit is contained in:
weixin_46229132 2025-03-14 09:45:46 +08:00
parent 64935bf92f
commit c1eb9d9528

View File

@ -7,12 +7,12 @@ from torch.distributions import Categorical
print("============================================================================================") print("============================================================================================")
# set device to cpu or cuda # set device to cpu or cuda
device = torch.device('cpu') device = torch.device('cpu')
if(torch.cuda.is_available()): # if(torch.cuda.is_available()):
device = torch.device('cuda:0') # device = torch.device('cuda:0')
torch.cuda.empty_cache() # torch.cuda.empty_cache()
print("Device set to : " + str(torch.cuda.get_device_name(device))) # print("Device set to : " + str(torch.cuda.get_device_name(device)))
else: # else:
print("Device set to : cpu") # print("Device set to : cpu")
print("============================================================================================") print("============================================================================================")