# Train GPT2 from scratch (Tensorflow) - Loss function issue **URL:** \ **Category:** Beginners\ **Created:** [March 11, 2021, 3:52pm UTC](https://discuss.huggingface.co/t/train-gpt2-from-scratch-tensorflow-loss-function-issue/4314 "2021-03-11T15:52:42Z")\ **Posts on this page:** 1\ **Page:** 1
**Author:** ![aldrodriguezca](https://avatars.discourse-cdn.com/v4/letter/a/848f3c/32.png) [@aldrodriguezca](https://discuss.huggingface.co/u/aldrodriguezca)\ **Post date:** [March 11, 2021, 3:52pm UTC](https://discuss.huggingface.co/t/train-gpt2-from-scratch-tensorflow-loss-function-issue/4314/1 "2021-03-11T15:52:42Z")
Hello. I’m trying to train a GPT2 model (actually GPT2LMHeadModel) using tensorflow2. I found [this post](https://towardsdatascience.com/train-gpt-2-in-your-own-language-fc6ad4d60171) where the author shows how to do it with great detail. However, there’s a special aspect regarding the definition of the loss function when compiling the model. Initially, when I was trying to implement it on my own, I defined the loss function as usual for the Keras models: ``` model.compile(optimizer=optimizer, loss=loss_fn, metrics=[metric]) ``` However, when trying to execute the `fit` method, it threw an error: ``` history = model.fit(dataset, epochs=EPOCHS) ValueError: Shape mismatch: The shape of labels (received (490,)) should equal the shape of logits except for the last dimension (received (11760, 64)) ``` In this case, the dateset corresponds to a tensorflow dataset, which shape is `` (Batch size: 10, sequences length: 49). After that, I realized the loss function parameter is defined differently in the tutorial mentioned above. The loss parameter corresponds to a list, which actually happens to be a function followed by a collection of None values, depending on the number of layers defined for the model architecture. Under this approach, it works fine. ![image](https://us1.discourse-cdn.com/hellohellohello/original/2X/4/44e8f046b02431dd6a2a3ed7b690c761b347694a.png) In both cases, the loss function corresponds to `tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)`. The documentation for `TFGPT2LMHeadModel` specifies: “`The GPT2 Model transformer with a language modeling head on top (linear layer with weights tied to the input embeddings).`”. So, I think it makes sense to define the loss as that list-based parameter (Basically the loss function would work for the top layer). After seeing this, I have some questions: - Is it correct to define the loss function that way? - Does it have any implications for the inference process? Additionally, I would like to mention that I’ve also tried to train the model using the [Trainer class](https://huggingface.co/transformers/main_classes/trainer.html), unfortunately, it throws a similar error when running the `train` method. ``` training_args = TFTrainingArguments( output_dir='./results', # output directory num_train_epochs=EPOCHS, # total # of training epochs per_device_train_batch_size=10, # batch size per device during training per_device_eval_batch_size=10, # batch size for evaluation warmup_steps=500, # number of warmup steps for learning rate scheduler weight_decay=0.01, # strength of weight decay logging_dir='./logs', # directory for storing logs) trainer = TFTrainer(model=model, args=training_args, train_dataset=dataset, eval_dataset=dataset) trainer.train() ``` ValueError: Shapes (248040,) and (4410,) are incompatible I’m using transformers v 3.5.0 and tokenizers v 0.9.3 Sorry for the long post, thanks in advance for your help! 🙂 🤗