def test_attention_outputs(self):
        r"""
        Overriding the test_attention_outputs test as the attention outputs of Rwkv are different from other models
        it has a shape `batch_size, seq_len, hidden_size`.
        """
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
        config.return_dict = True

        seq_len = getattr(self.model_tester, "seq_length", None)

        for model_class in self.all_model_classes:
            inputs_dict["output_attentions"] = True
            inputs_dict["output_hidden_states"] = False
            config.return_dict = True
            model = model_class(config)
            model.to(torch_device)
            model.eval()

            inputs = self._prepare_for_class(inputs_dict, model_class)
            batch_size = inputs["input_ids"].shape[0]
            with torch.no_grad():
                outputs = model(**inputs)
            attentions = outputs.attentions
            self.assertEqual(len(attentions), self.model_tester.num_hidden_layers)

            # check that output_attentions also work using config
            del inputs_dict["output_attentions"]
            config.output_attentions = True
            model = model_class(config)
            model.to(torch_device)
            model.eval()

            inputs = self._prepare_for_class(inputs_dict, model_class)
            batch_size = inputs["input_ids"].shape[0]
            with torch.no_grad():
                outputs = model(**inputs)
            attentions = outputs.attentions
            self.assertEqual(len(attentions), self.model_tester.num_hidden_layers)

            self.assertListEqual(
                list(attentions[0].shape[-3:]),
                [batch_size, seq_len, config.hidden_size],
            )
            out_len = len(outputs)

            # Check attention is always last and order is fine
            inputs_dict["output_attentions"] = True
            inputs_dict["output_hidden_states"] = True
            model = model_class(config)
            model.to(torch_device)
            model.eval()

            inputs = self._prepare_for_class(inputs_dict, model_class)
            batch_size = inputs["input_ids"].shape[0]
            with torch.no_grad():
                outputs = model(**inputs)

            added_hidden_states = 1
            self.assertEqual(out_len + added_hidden_states, len(outputs))

            self_attentions = outputs.attentions

            self.assertEqual(len(self_attentions), self.model_tester.num_hidden_layers)
            self.assertListEqual(
                list(self_attentions[0].shape[-3:]),
                [batch_size, seq_len, config.hidden_size],
            )

        Overriding the test_attention_outputs test as the attention outputs of Rwkv are different from other models
        it has a shape `batch_size, seq_len, hidden_size`.
        
dpkd/transformers/tests/models/rwkv/test_modeling_rwkv.py