Reducing size of training dataset in tensorflow 2 tutorial (Transformer model for language understanding) with '.take(n)' method does not work

Question

I am a Tensorflow-newbie, therefore forgive me if my question is too basic or stupid ;)

I tried to reduce the size of the training dataset in the "Transformer mo