Problem in the Vertex AI example of 'Multi-class prediction with a DNN hosted on Vertex AI'

[Liulizhi]

I runed the Vertex AI example shown above in cloab, when it reaches the De/serialization stage, the error message is as follows.How should I solve this problem?
ValueError: Exception encountered when calling DeSerializeInput.call(). Could not automatically infer the output shape / dtype of 'de_serialize_input' (of type DeSerializeInput). Either the `DeSerializeInput.call()` method is incorrect, or you need to implement the `DeSerializeInput.compute_output_spec() / compute_output_shape()` method. Error encountered: as_list() is not defined on an unknown TensorShape. 

In addition, in another example of Hosting a Pre-trained Tree-crown Segmentation Model, a similar problem also occurred, but I was able to run it normally in November 2023.
NotImplementedError: Exception encountered when calling Lambda.call(). We could not automatically infer the shape of the Lambda's output. Please specify the `output_shape` argument for this Lambda layer.