def METHOD_NAME(self):
    model = DummyModule().cuda()
    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
    state = deepcopy(model.state_dict())
    config = EMAConfig()
    ema = EMA(model, config)
    # set decay
    ema._set_decay(config.ema_decay)
    self.assertEqual(ema.get_decay(), config.ema_decay)
    # get model
    self.assertEqual(ema.get_model(), ema.model)
    # Since fp32 params is not used, it should be of size 0
    self.assertEqual(len(ema.fp32_params), 0)
    # EMA step
    x = torch.randn(32).cuda()
    y = model(x)
    loss = y.sum()
    loss.backward()
    optimizer.step()
    ema.step(model)
    ema_state_dict = ema.get_model().state_dict()
    for key, param in model.state_dict().items():
        prev_param = state[key]
        ema_param = ema_state_dict[key]
        if "version" in key:
            # Do not decay a model.version pytorch param
            continue
        self.assertTorchAllClose(
            ema_param,
            config.ema_decay * prev_param + (1 - config.ema_decay) * param,
        )
    # Since fp32 params is not used, it should be of size 0
    self.assertEqual(len(ema.fp32_params), 0)
    # Load EMA into model
    model2 = DummyModule().cuda()
    ema.reverse(model2)
    for key, param in model2.state_dict().items():
        ema_param = ema_state_dict[key]
        self.assertTrue(torch.allclose(ema_param, param))