def METHOD_NAME(self, input_param, input_shape, _):
    net = ViT(**(input_param))
    net.eval()
    with torch.no_grad():
        torch.jit.script(net)
    test_data = torch.randn(input_shape)
    test_script_save(net, test_data)