code: def test_dialogue_nearest_neighbour_mean_pooling():
      model_output = [torch.ones(8, 512, 768)]
      attention_mask = torch.ones(8, 512)
      assert torch.equal(
          torch.ones(8, 768).float(),
          DialogueNearestNeighbourModel.mean_pooling(
              model_output, attention_mask
          ),
      )
  
      model_output = [torch.zeros(8, 512, 768)]
      attention_mask = torch.ones(8, 512)
      assert torch.equal(
          torch.zeros(8, 768).float(),
          DialogueNearestNeighbourModel.mean_pooling(
              model_output, attention_mask
          ),
      )
  
      model_output = [
          torch.cat(
              [torch.zeros(8, 256, 768), torch.ones(8, 256, 768)],
              axis=1,
          )
      ]
      attention_mask = torch.ones(8, 512)
      assert torch.equal(
          torch.ones(8, 768).float() * 0.5,
          DialogueNearestNeighbourModel.mean_pooling(
              model_output, attention_mask
          ),
      )
  
imports: import pytest
  import torch
  from nemo.collections.nlp.data.dialogue.data_processor.assistant_data_processor import DialogueAssistantDataProcessor
  from nemo.collections.nlp.data.dialogue.data_processor.data_processor import DialogueDataProcessor
  from nemo.collections.nlp.data.dialogue.data_processor.sgd_data_processor import DialogueSGDDataProcessor
  from nemo.collections.nlp.data.dialogue.dataset.dialogue_gpt_classification_dataset import DialogueGPTClassificationDataset
  from nemo.collections.nlp.data.dialogue.dataset.dialogue_s2s_generation_dataset import DialogueS2SGenerationDataset
  from nemo.collections.nlp.data.dialogue.dataset.dialogue_sgd_bert_dataset import DialogueSGDBERTDataset
  from nemo.collections.nlp.metrics.dialogue_metrics import DialogueClassificationMetrics, DialogueGenerationMetrics
  from nemo.collections.nlp.models.dialogue.dialogue_nearest_neighbour_model import DialogueNearestNeighbourModel
label_desc: Use image node_id 12 for calling a global function with example usage: test_dialogue_nearest_neighbour_mean_pooling() without return types