Trainer use one epoch of the test_dataloader

I want to try the test mode with trainer.test(model, data_module, ckpt_path="./checkpoints/best-checkpoint.ckpt"). If I did not make any mistake, with that line of cod, the model will work with the dataset_test that I created in the LightningDataModule and run each epoch, calculate the loss and at the end it will return el avg_loss (code below).

class Data(pl.LightningDataModule): #data_module
  def __init__(self, test_df):
    self.test_df = test_df

  def setup(self, stage=None):
    self.dataset_test = TensorDataset(input_ids_test, attention_masks_test, labels_test)

    def test_dataloader(self):
    return DataLoader(self.dataset_test,
class ClassModel(pl.LightningModule): #model
  def __init__(self):
    self.model = bert_model

  def configure_optimizers(self):

  def forward(self, input_ids, attention_mask, labels):
    loss, output = self.model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
    return loss, output


  def test_step(self, batch, batch_idx):
    inputs = {'input_ids':      batch[0],
              'attention_mask': batch[1],
              'labels':         batch[2]}
    outputs = self.model(**inputs)
    loss = outputs[0]
    self.log("Test loss", loss, on_epoch=True)
    return {'loss': loss}

  def test_epoch_end(self, outputs):
    avg_loss = torch.stack([x['loss'] for x in outputs]).mean()
    return {'avg_loss': avg_loss}

The result that I got after I run trainer.test(model, data_module) is :

We can see her that it only run for one epoch (63) while the size of my test dataset is 1000 and the batch size is 16.
Can you please tell me here if I missed something?
Thank you for your time!

hey @BttMA

max_epochs doesn’t have any effect on testing/validation/prediction. It just iterates over the complete dataset only once. Since there is no optimization happening here, running .test multiple times won’t provide any significance. I hope it answers your question.

Also, we have moved the discussions to GitHub Discussions. You might want to check that out instead to get a quick response. The forums will be marked read-only after some time.

Thank you

Hello :slight_smile: Thank you for your answer. It has been a while since I posted this qu. and I think I didn’t make myself very clear as I mixed up between epoch and batch (sorry!!). what I wanted to test is all the test dataset not only one batch of it. I saw in the doc that the trainer.test run only over one epoch but you didn’t mention the batch size. The snippet of my code show that the test run only over one batch and not over all the test dataset. I hope this will be more clear. Thanks :wink:

this looks good to me. since your total dataset size is 1000 and batch_size is 16, total batches = 1000/16 ~= 63, and it shows 63 as the final counter.

:face_with_peeking_eye: it looks fine to me now :sweat_smile:
Maybe, back then, I misread my data.
Thanks a lot