Cacher is the code snippet organizer for pro developers

import torch
def get_device(config):
    # 检查GPU是否可用
    use_cuda = torch.cuda.is_available()
    # 　如果可用使用0号GPU
    device = torch.device("cuda:0" if use_cuda else "cpu")
    if config['custom']['use_cpu'] is True:
        device = torch.device('cpu')
    config['custom']['device'] = device
    return device