If someone else in the future has the exact same error - the OP had a mistake in Exercise 1 - Encoder:
# Pass the context through the embedding layer
x = self.embedding(tf.keras.Input(tensor=context))
This would pass unit tests but would cause the error during training since this is not correct way of doing the embedding.
We use the call method to define the forward pass of our layer, and we would expect the input to be passed as an argument to the call method, not created within it.
So the correct way of implementing the embedding should be:
# Pass the context through the embedding layer
x = self.embedding(context)
Cheers