Machine learning neophyte here, so apologies in advance for a “dumb” question.
I have been trying to build a TrOCR model using the VisionEncoderDecoderModel with a checkpoint ‘microsoft/trocr-base-handwritten’ . I have tried the other ones too, but my fine tuning messes up the model instead of improving it.  Wanted to ask for some help with fixing this/understanding what goes wrong. I have been using pytorch lightning for the training/fine tuning. My code is below. Out of the box (with the above checkpoint) model can generate pretty accurate results, but after my training/fine tuning its gets worse instead of better.
Some info: I am fine tuning on IAM dataset. The initial loss, when starting is around 8 and it never goes below 4…
Huge thanks for any help!
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
processor = TrOCRProcessor.from_pretrained('microsoft/trocr-base-handwritten')

class TrOCR_Image_to_Text(pl.LightningModule):
    def __init__(self):
        super().__init__()
        model = VisionEncoderDecoderModel.from_pretrained('microsoft/trocr-base-handwritten') 
        model.config.decoder_start_token_id = processor.tokenizer.cls_token_id
        model.config.pad_token_id = processor.tokenizer.pad_token_id
        model.config.vocab_size = model.config.decoder.vocab_size

        # set beam search parameters
        model.config.eos_token_id = processor.tokenizer.sep_token_id
        model.config.max_length = 89
        model.config.early_stopping = True
        model.config.no_repeat_ngram_size = 3
        model.config.length_penalty = 2.0
        model.config.num_beams = 2

        self.vit = model

    def generate(self, input_ids):
        return self.vit.generate(input_ids)

    def forward(self, batch):
        model = self.vit
        model.to(device)
        x,y = batch
        pixel_values = x
        labels = y

        outputs = model(pixel_values=pixel_values, labels=labels)
        loss = outputs.loss
        logits = outputs.logits

        return loss

    def training_step(self, batch, batch_idx):
        loss = self.forward(batch)
        self.log("training_loss", loss)

        return loss


    def validation_step(self, batch, batch_idx):
        loss = self.forward(batch)
        self.log("validation_loss", loss, prog_bar=True, on_epoch=True)

        model = self.vit
        
        x,y = batch

        outputs = model.generate(x.to(device))

        cer = compute_cer(pred_ids=outputs, label_ids=y)
        valid_cer = cer

        self.log("cer", valid_cer, prog_bar=True)

        return loss

    def test_step(self, batch, batch_idx):
        loss = self(batch, batch_idx)

        return loss

 def configure_optimizers(self):
        return AdamW(self.parameters(), lr=5e-5)

    def train_dataloader(self):
        return train_dataloader

    def val_dataloader(self):
        return eval_dataloader

    def test_dataloader(self):
        return eval_dataloader


'''
CER Metric:

'''
from datasets import load_metric
cer_metric = load_metric("cer")

def compute_cer(pred_ids, label_ids):
    pred_str = processor.batch_decode(pred_ids, skip_special_tokens=True)
    label_ids[label_ids == -100] = processor.tokenizer.pad_token_id
    label_str = processor.batch_decode(label_ids, skip_special_tokens=True)

    cer = cer_metric.compute(predictions=pred_str, references=label_str)

    return cer