Function save_model

examples/pytorch/vit/ViT-quantization/eval_engine.py:91–95

(args, model)

def save_model(args, model):
    model_to_save = model.module if hasattr(model, 'module') else model
    model_checkpoint = os.path.join(args.output_dir, "%s_checkpoint.bin" % args.name)
    torch.save(model_to_save.state_dict(), model_checkpoint)
    logger.info("Saved model checkpoint to %s", model_checkpoint)


def setup(args):