code: def test_averaged_model_mixed_device(self, ema):
      if not torch.cuda.is_available():
          return
      dnn = torch.nn.Sequential(
          torch.nn.Conv2d(1, 5, kernel_size=3), torch.nn.Linear(5, 10)
      )
      dnn[0].cuda()
      dnn[1].cpu()
  
      averaged_params, averaged_dnn = self._run_averaged_steps(
          dnn, None, ema
      )
  
      for p_avg, p_swa in zip(
          averaged_params, averaged_dnn.parameters()
      ):
          self.assertEqual(p_avg, p_swa)
          # Check that AveragedModel is on the correct device
          self.assertTrue(p_avg.device == p_swa.device)
  
file_path: pytorch/test/optim/test_swa_utils.py
imports: import itertools
  import pickle
  import torch
  from torch.optim.swa_utils import AveragedModel, update_bn, get_swa_multi_avg_fn, get_ema_multi_avg_fn
  from torch.testing._internal.common_utils import TestCase, load_tests, parametrize, instantiate_parametrized_tests
label_desc: Use image node_id 4 for calling the TestSWAUtils obj's underlying member method code with example usage: obj.test_averaged_model_mixed_device(ema) without return types