def METHOD_NAME(self):
    optimizer_mock = MagicMock()
    lr_scheduler_mock = MagicMock()
    prepared_optimizer, prepared_scheduler = self.elastic_trainer.prepare(
        optimizer_mock, lr_scheduler_mock
    )
    self.assertIsInstance(prepared_optimizer, _ElasticOptimizer)
    self.assertIsInstance(prepared_scheduler, _ElasticLRScheduler)