def METHOD_NAME(self):
    model = DummyModule().half()
    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
    state = deepcopy(model.state_dict())
    config = EMAConfig(ema_fp32=False)
    ema = EMA(model, config)
    # Since fp32 params is not used, it should be of size 0
    self.assertEqual(len(ema.fp32_params), 0)
    x = torch.randn(32)
    y = model(x.half())
    loss = y.sum()
    loss.backward()
    optimizer.step()
    ema.step(model)
    for key, param in model.state_dict().items():
        prev_param = state[key]
        ema_param = ema.get_model().state_dict()[key]
        if "version" in key:
            # Do not decay a model.version pytorch param
            continue
        # EMA update is done in fp16, and hence the EMA param must be
        # closer to the EMA update done in fp16 than in fp32.
        self.assertLessEqual(
            torch.norm(
                ema_param.float() -
                (config.ema_decay * prev_param + (1 - config.ema_decay) * param).float()
            ),
            torch.norm(
                ema_param.float() -
                (config.ema_decay * prev_param.float() + (1 - config.ema_decay) * param.float()).half().float()
            ),
        )
        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)