cpu to device select

This commit is contained in:
王庆刚
2025-01-13 18:11:56 +08:00
parent a50f777839
commit 744fb7b7b2
4 changed files with 3 additions and 3 deletions

View File

@ -76,8 +76,8 @@ class Config:
lr_decay = 0.95 # 0.98
weight_decay = 5e-4
loss = 'cross_entropy' # ['focal_loss', 'cross_entropy']
device = torch.device('cuda:1' if torch.cuda.is_available() else 'cpu')
# device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
# device = torch.device('cuda:1' if torch.cuda.is_available() else 'cpu')
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
pin_memory = True # if memory is large, set it True to speed up a bit
num_workers = 4 # dataloader