Qianli Scott Zhu

Sorry for the long wait, since end user could easily miss the preprocess API with pytorch format, how about we include the preprocess as part of the model, and control...

@k-w-w, can you take a look for this PR? Thanks.

Assign to Rick for further inputs.

Ping @rchao for review.

(Removing the label is @rchao is the active reviewer for this).

Took a look for the PR and the repro, it is bit weird that the layer and model are treating and tracking the weights differently. From API perspective, the subclass...

Adding @rohan100jain from tf.core team for this issue. From the latest message, it might be a int32 overflow issue somewhere in TF. In the meantime, please provide a some reproducible...

I think the issue here is that the model output shape is [404, 1] and the label shape is [404]. The MSE produce the incorrect value since the label and...

I think this is tricky error and an easy pitfall for end user. We should either error out when there is a shape mismatch, or broadcast to same shape.

Thanks for the detailed bug report, seems to be a quite important issue since MHA is core API.