code: def test_batch_size_per_device(self, a):
      # Need to patch the mmf.utils.general's world size not mmf.utils.distributed
      # as the first one is what will be used
      with patch("mmf.utils.general.get_world_size", return_value=2):
          config = self._get_config(
              max_updates=2, max_epochs=None, batch_size=4
          )
          trainer = TrainerTrainingLoopMock(config=config)
          add_model(trainer, SimpleModel({"in_dim": 1}))
          add_optimizer(trainer, config)
          registry.register("config", trainer.config)
          batch_size = get_batch_size()
          trainer.config.training.batch_size = batch_size
          trainer.load_datasets()
          # Train loader has batch size per device, for global batch size 4
          # with world size 2, batch size per device should 4 // 2 = 2
          self.assertEqual(
              trainer.train_loader.current_loader.batch_size, 2
          )
          # This is per device, so should stay same
          config = self._get_config(
              max_updates=2, max_epochs=None, batch_size_per_device=4
          )
          trainer = TrainerTrainingLoopMock(config=config)
          add_model(trainer, SimpleModel({"in_dim": 1}))
          add_optimizer(trainer, config)
          registry.register("config", trainer.config)
          batch_size = get_batch_size()
          trainer.config.training.batch_size = batch_size
          trainer.load_datasets()
          self.assertEqual(
              trainer.train_loader.current_loader.batch_size, 4
          )
  
      max_updates = trainer._calculate_max_updates()
      self.assertEqual(max_updates, 2)
  
      self.check_values(trainer, 0, 0, 0)
      trainer.training_loop()
      self.check_values(trainer, 2, 1, 2)
  
imports: import unittest
  from unittest.mock import MagicMock, patch
  import torch
  from mmf.common.registry import registry
  from mmf.utils.general import get_batch_size
  from tests.test_utils import SimpleModel, SimpleNaNLossModel
  from tests.trainers.test_trainer_mocks import TrainerTrainingLoopMock
  from tests.trainers.test_utils import add_model, add_optimizer, get_config_with_defaults, get_mmf_trainer
label_desc: Use image node_id 11 for calling the TestTrainingLoop obj's underlying member method code with example usage: obj.test_batch_size_per_device(a) without return types