code: def __init__(
      self,
      config,
      num_train_data,
      max_updates,
      max_epochs,
      device="cuda",
  ):
      config.training.max_updates = max_updates
      config.training.max_epochs = max_epochs
      config.training.fp16 = True
      config = get_config_with_defaults(config)
  
      super().__init__(num_train_data, config=config)
      self.device = torch.device(device)
      self.config = config
      self.model = SimpleModel({"in_dim": 1})
      self.model.build()
      self.model = self.model.cuda()
      self.optimizer = build_optimizer(self.model, self.config)
      self.distributed = True
      self.local_rank = 0
      self.parallelize_model()
      self.load_fp16_scaler()
  
file_path: mmf/tests/trainers/test_sharded_ddp.py
imports: import os
  import unittest
  import torch
  from mmf.common.registry import registry
  from mmf.trainers.mmf_trainer import MMFTrainer
  from mmf.utils.build import build_optimizer
  from omegaconf import OmegaConf
  from tests.test_utils import SimpleModel, skip_if_no_cuda
  from tests.trainers.test_training_loop import TrainerTrainingLoopMock
  from tests.trainers.test_utils import get_config_with_defaults
label_desc: Use image node_id 1 to create a new MMFTrainerMock object from inherited base classes: TrainerTrainingLoopMock, MMFTrainer with example: obj = MMFTrainerMock(config, num_train_data, max_updates, max_epochs, device)