code: def process_batch(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
  
file_path: determined/examples/features/torch_batch_process_embeddings/bert_embedding_generation.py
imports: import logging
  import os
  import pathlib
  import shutil
  import chromadb
  import torch
  from datasets import load_dataset
  from transformers import BertModel, BertTokenizer
  from determined.pytorch import experimental
label_desc: Use image node_id 2 for calling the EmbeddingProcessor obj's underlying member method code with example usage: obj.process_batch(batch, batch_idx) without return types