def build_model(): model = models.__dict__[args.arch](pretrained=True, num_classes=101) model.cuda() return model