code: def test_loss_container(num_batches):
      """Test `LossContainer` wrapper for classification task."""
      loss = nn.CrossEntropyLoss(reduction="mean")
      loss_container = LossContainer(loss=loss, name="CrossEntropyLoss")
      cumulated_loss_val, counts = 0, 0
  
      for i in range(num_batches):
          input = torch.randn(3, 5)
          target = torch.empty(3, dtype=torch.long).random_(5)
  
          loss_val1 = loss_container(y_pred=input, y_true=target)
          loss_val2 = loss(input, target)
          assert torch.allclose(loss_val1, loss_val2)
  
          cumulated_loss_val += loss_val2.item()
          counts += 1
  
      expected_val = cumulated_loss_val / counts
      result = loss_container.result()
      assert np.isclose(expected_val, result["CrossEntropyLoss"])
  
      loss_container.reset()
      assert np.isclose(loss_container.total, 0)
      assert np.isclose(loss_container.count, 0)
  
imports: import numpy
  import pytest
  import torch
  import torch.nn
  from alibi.models.pytorch.metrics import AccuracyMetric, LossContainer
label_desc: Use image node_id 1 for calling a global function with example usage: test_loss_container(num_batches) without return types