code: def test_ema(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))
  
file_path: fairseq/tests/gpu/test_ema_gpu.py
imports: import unittest
  from copy import deepcopy
  from dataclasses import dataclass
  from typing import Optional
  import torch
  from fairseq.models.ema import EMA
label_desc: Use image node_id 2 for calling the TestEMAGPU obj's underlying member method code with example usage: obj.test_ema() without return types