docstr: Let's make sure we can actually compute the loss and do a backward on it.
docstr_tok: ["Let","'s","make","sure","we","can","actually","compute","the","loss","and","do","a","backward","on","it","."]
code: def test_load_balancing_loss(self):
      r"""
      Let's make sure we can actually compute the loss and do a backward on it.
      """
  
      (
          config,
          input_dict,
      ) = self.model_tester.prepare_config_and_inputs_for_common()
      config.num_labels = 3
      config.num_local_experts = 8
      config.output_router_logits = True
      input_ids = input_dict["input_ids"]
      attention_mask = input_ids.ne(1).to(torch_device)
      model = MixtralForCausalLM(config)
      model.to(torch_device)
      model.eval()
      result = model(input_ids, attention_mask=attention_mask)
      self.assertEqual(
          result.router_logits[0].shape, (91, config.num_local_experts)
      )
      torch.testing.assert_close(
          result.aux_loss.cpu(),
          torch.tensor(2, dtype=torch.float32),
          rtol=1e-2,
          atol=1e-2,
      )
  
imports: import tempfile
  import unittest
  import pytest
  from transformers import MixtralConfig, is_torch_available
  from transformers.testing_utils import require_flash_attn, require_torch, require_torch_gpu, slow, torch_device
  from ...generation.test_utils import GenerationTesterMixin
  from ...test_configuration_common import ConfigTester
  from ...test_modeling_common import ModelTesterMixin, ids_tensor
  from ...test_pipeline_mixin import PipelineTesterMixin
label_desc: Use image node_id 14 for calling the MixtralModelTest obj's underlying member method code with example usage: obj.test_load_balancing_loss() without return types