ValueError: Dimension size must be evenly divisible by 250 but is 1 for '{{node Reshape}}

Hello Community,

I'm currently working on a project which requires me to include differential privacy. While I try to implement it using tf privacy, I become an error message:

ValueError: Dimension size must be evenly divisible by 250 but is 1 for '{{node Reshape}} = Reshape[T=DT_FLOAT, Tshape=DT_INT32](sparse_categorical_crossentropy/weighted_loss/value, Reshape/shape)' with input shapes: [], [2] and with input tensors computed as partial shapes: input[1] = [250,?].

The error takes place in line 381 where the model should be fitted. Because my code is over 400 lines long and I have no clue which lines are interesting, I will share it with you via codeshare: https://codeshare.io/mpvkVj

If the issue gets resolved, I will include the interesting lines later on. Thank you very much for every hint 😃

Start with the model-fitting call at line 381 and inspect the shared Codeshare example alongside the reported Reshape error. Trace the tensors involved in the weighted loss and identify the relevant minimal lines; done means the cause is reproducible and clearly explained.