code: def test_transformer_jit_embeddings(self, results_base_path):
      base_embeddings = TransformerWordEmbeddings(
          "distilbert-base-uncased",
          layers="-1,-2,-3,-4",
          layer_mean=False,
          allow_long_sentences=True,
      )
      sentence: Sentence = Sentence(
          "I love Berlin, but Vienna is where my hearth is."
      )
  
      class JitWrapper(torch.nn.Module):
          def __init__(
              self, embedding: TransformerWordEmbeddings
          ) -> None:
              super().__init__()
              self.embedding = embedding
  
          def forward(
              self,
              input_ids: torch.Tensor,
              token_lengths: torch.LongTensor,
              attention_mask: torch.Tensor,
              overflow_to_sample_mapping: torch.Tensor,
              word_ids: torch.Tensor,
          ):
              return self.embedding.forward(
                  input_ids=input_ids,
                  token_lengths=token_lengths,
                  attention_mask=attention_mask,
                  overflow_to_sample_mapping=overflow_to_sample_mapping,
                  word_ids=word_ids,
              )["token_embeddings"]
  
      base_embeddings.embed(sentence)
      base_token_embedding = sentence[5].get_embedding().clone()
      sentence.clear_embeddings()
  
      tensors = base_embeddings.prepare_tensors([sentence])
      # ensure that the prepared tensors is what we expect
      assert sorted(tensors.keys()) == [
          "attention_mask",
          "input_ids",
          "overflow_to_sample_mapping",
          "token_lengths",
          "word_ids",
      ]
  
      wrapper = JitWrapper(base_embeddings)
      (
          parameter_names,
          parameter_list,
      ) = TransformerJitWordEmbeddings.parameter_to_list(
          base_embeddings, wrapper, [sentence]
      )
      with warnings.catch_warnings():
          warnings.simplefilter("ignore")
          script_module = torch.jit.trace(wrapper, parameter_list)
      jit_embeddings = (
          TransformerJitWordEmbeddings.create_from_embedding(
              script_module, base_embeddings, parameter_names
          )
      )
  
      jit_embeddings.embed(sentence)
      jit_token_embedding = sentence[5].get_embedding().clone()
      assert torch.isclose(
          base_token_embedding, jit_token_embedding
      ).all()
      sentence.clear_embeddings()
  
      # use a SequenceTagger to save and reload the embedding in the manner it is supposed to work
      example_tagger = SequenceTagger(
          embeddings=jit_embeddings,
          tag_dictionary=Dictionary(),
          tag_type="none",
      )
      results_base_path.mkdir(exist_ok=True, parents=True)
      example_tagger.save(results_base_path / "tagger.pt")
      del example_tagger
      new_example_tagger = SequenceTagger.load(
          results_base_path / "tagger.pt"
      )
      loaded_jit_embedding = new_example_tagger.embeddings
  
      loaded_jit_embedding.embed(sentence)
      loaded_jit_token_embedding = sentence[5].get_embedding().clone()
      sentence.clear_embeddings()
      assert torch.isclose(
          jit_token_embedding, loaded_jit_token_embedding
      ).all()
  
imports: import importlib.util
  import warnings
  import pytest
  import torch
  from PIL import Image
  from transformers.utils import is_detectron2_available
  from flair.data import BoundingBox, Dictionary, Sentence
  from flair.embeddings import TransformerJitWordEmbeddings, TransformerWordEmbeddings
  from flair.models import SequenceTagger
  from tests.embedding_test_utils import BaseEmbeddingsTest
label_desc: Use image node_id 1 for calling the TestTransformerWordEmbeddings obj's underlying member method code with example usage: obj.test_transformer_jit_embeddings(results_base_path) and returns: self