tfm.nlp.layers.RelativePositionEmbedding

[Leo-Lifeblood]

Issue type

Bug

Have you reproduced the bug with TensorFlow Nightly?

No

Source

source

TensorFlow version

2.16.1

Custom code

Yes

OS platform and distribution

google colab

Mobile device

No response

Python version

No response

Bazel version

No response

GCC/compiler version

No response

CUDA/cuDNN version

No response

GPU model and memory

No response

Current behavior?

Standalone code to reproduce the issue

import tensorflow as tf
import tensorflow_models as tfm

input_tensor = tf.random.normal((32, 87, 128))

position_emb = tfm.nlp.layers.RelativePositionEmbedding(hidden_size=128)

output_tensor = position_emb(input_tensor)

print(output_tensor.shape)

Relevant log output

No response

[sushreebarsa]

@Leo-Lifeblood I was able to replicate the output reported here, could you please share the error you are encountering ? Thank you!

[Leo-Lifeblood]

The error is that the batch dimension in this case 32 is being deleted. This breaks any training loop.