cpu to device select
This commit is contained in:
@ -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
|
||||
|
Reference in New Issue
Block a user