code: def test_logits_with_cache(
      self, test_name: str, model_type: str, batch_size: int
  ):
      input_ids = torch.randint(low=1, high=10, size=(batch_size, 1))
      seq_length = 12
      attention_mask = torch.ones(
          batch_size, seq_length + 1, dtype=torch.int32
      )
  
      model_id = MODELS_DICT[model_type]
  
      model = AutoModelForCausalLM.from_pretrained(model_id)
  
      normalized_config = (
          NormalizedConfigManager.get_normalized_config_class(
              model.config.model_type
          )(model.config)
      )
  
      if model_type == "bloom":
          pkv_generator_class = BloomDummyPastKeyValuesGenerator
      else:
          pkv_generator_class = DummyPastKeyValuesGenerator
  
      pkv_generator = pkv_generator_class(
          task="",
          normalized_config=normalized_config,
          batch_size=batch_size,
          sequence_length=seq_length,
      )
  
      past_key_values = pkv_generator.generate(
          input_name="past_key_values"
      )
  
      result_vanilla = model(
          input_ids=input_ids,
          attention_mask=attention_mask,
          past_key_values=past_key_values,
      )
  
      model = BetterTransformer.transform(model)
      result_bettertransformer = model(
          input_ids=input_ids,
          attention_mask=attention_mask,
          past_key_values=past_key_values,
      )
  
      logits_vanilla = result_vanilla.logits
      logits_bettertransformer = result_bettertransformer.logits
      self.assertTrue(
          torch.allclose(
              logits_vanilla, logits_bettertransformer, atol=1e-5
          ),
          f" Maxdiff: {(logits_vanilla - logits_bettertransformer).abs().max()}",
      )
  
imports: import gc
  import unittest
  import pytest
  import torch
  from parameterized import parameterized
  from testing_utils import MODELS_DICT, BetterTransformersTestMixin
  from transformers import AutoModelForCausalLM, AutoTokenizer
  from optimum.bettertransformer import BetterTransformer
  from optimum.utils import BloomDummyPastKeyValuesGenerator, DummyPastKeyValuesGenerator, NormalizedConfigManager
  from optimum.utils.testing_utils import grid_parameters, require_accelerate, require_torch_gpu
label_desc: Use image node_id 5 for calling the BetterTransformersDecoderTest obj's underlying member method code with example usage: obj.test_logits_with_cache(test_name, model_type, batch_size) without return types