def METHOD_NAME(self, batch, batch_idx) -> None:
    with torch.no_grad():
        tokenized_text = self.tokenizer.batch_encode_plus(
            batch["text"],
            truncation=True,
            padding="max_length",
            max_length=512,
            add_special_tokens=True,
        )
        inputs = torch.tensor(tokenized_text["input_ids"])
        inputs = inputs.to(self.device)
        masks = torch.tensor(tokenized_text["attention_mask"])
        masks = masks.to(self.device)
        outputs = self.model(inputs, masks)
        # To create an embedding vector for each document,
        # 1. we take the hidden states from the last layer (output["hidden_states"][-1]),
        #    which is a tensor of (#examples, #tokens, #hidden_states) size.
        # 2. we calculate the average across the token-dimension, resulting in a tensor of
        #    (#examples, #hidden_states) size.
        outputs = torch.mean(outputs["hidden_states"][-1], dim=1)
        self.output.append({"embeddings": outputs, "id": batch["_id"], "text": batch["text"]})
        self.last_index = batch_idx