def build_model(args): model = DeepGL(args.num_blocks) if args.pretrained_path: model.load_state_dict(torch.load( os.path.join(args.pretrained_path, 'samples') + '/' + str(args.load_step) + '.pt')) return model