INNER CODE UNIT · Python

main

pprp/SimpleCVReproduction · NAS/AngleNAS/DARTS/shrinking/main.py:48

def main():
    if not torch.cuda.is_available():
        print('no gpu device available')
        sys.exit(1)

    num_gpus = torch.cuda.device_count() 
    np.random.seed(args.seed)
    args.gpu = args.local_rank % num_gpus
    torch.cuda.set_device(args.gpu)
    cudnn.benchmark = True
    cudnn.deterministic = True
    torch.manual_seed(args.seed)
    cudnn.enabled=True
    torch.cuda.manual_seed(args.seed)
    group_name = 'search_space_shrinking'
    print('gpu device = %d' % args.gpu)
    print("args = %s", args)

View source record →