code: def test_parallelize_bmuf(self, cpu_device):
      assume(
          cpu_device
          or workspace.has_gpu_support
          or workspace.has_hip_support
      )
  
      workspace.ResetWorkspace()
  
      model = cnn.CNNModelHelper(order="NHWC", name="test")
      devices = [0, 1]
  
      def input_builder_fun(model):
          return None
  
      if not cpu_device:
          device_type = workspace.GpuDeviceType
          device_prefix = "gpu"
      else:
          device_type = caffe2_pb2.CPU
          device_prefix = "cpu"
      self._generate_data(devices, device_type, device_prefix)
  
      data_parallel_model.Parallelize_BMUF(
          model,
          input_builder_fun,
          self._model_build_fun,
          self._param_update_fun,
          devices=devices,
          cpu_device=cpu_device,
      )
  
      data_parallel_model.RunInitNet(model)
  
      # Check initial momentum params are zeros
      self.assertEqual(
          list(model._device_grouped_blobs.keys()), ["fc_w", "fc_b"]
      )
      self.assertEqual(
          workspace.FetchBlob("{}_0/fc_b_v".format(device_prefix)), 0
      )
      np.testing.assert_equal(
          workspace.FetchBlob("{}_0/fc_w_v".format(device_prefix)),
          np.zeros(16).astype(np.float32).reshape(1, 16),
      )
  
      # Run the algorithm for one iteration to have non-zero params.
      data_parallel_model.RunNet(model, 1)
  
      # Save iteration momentum and post local update params
      v_b_ = workspace.FetchBlob("{}_0/fc_b_v".format(device_prefix))
      v_w_ = workspace.FetchBlob("{}_0/fc_w_v".format(device_prefix))
  
      workspace.RunNetOnce(model.net)
  
      b_0_ = workspace.FetchBlob("{}_0/fc_b".format(device_prefix))
      w_0_ = workspace.FetchBlob("{}_0/fc_w".format(device_prefix))
      b_1_ = workspace.FetchBlob("{}_1/fc_b".format(device_prefix))
      w_1_ = workspace.FetchBlob("{}_1/fc_w".format(device_prefix))
  
      # Compute block gradients.
      b_g_ = workspace.FetchBlob("{}_0/fc_b_g".format(device_prefix))
      w_g_ = workspace.FetchBlob("{}_0/fc_w_g".format(device_prefix))
      workspace.RunNetOnce(model._global_model_param_updates_net)
  
      g_b = (b_0_ + b_1_) / 2 - b_g_
      g_w = (w_0_ + w_1_) / 2 - w_g_
      v_b = workspace.FetchBlob("{}_0/fc_b_v".format(device_prefix))
      v_w = workspace.FetchBlob("{}_0/fc_w_v".format(device_prefix))
  
      w_g = workspace.FetchBlob("{}_0/fc_w_g".format(device_prefix))
      b_g = workspace.FetchBlob("{}_0/fc_b_g".format(device_prefix))
      w_0 = workspace.FetchBlob("{}_0/fc_w".format(device_prefix))
      b_0 = workspace.FetchBlob("{}_0/fc_b".format(device_prefix))
      w_1 = workspace.FetchBlob("{}_1/fc_w".format(device_prefix))
      b_1 = workspace.FetchBlob("{}_1/fc_b".format(device_prefix))
  
      # Check momentum update step
      np.testing.assert_equal(v_b, 0.5 * v_b_ + g_b)
      np.testing.assert_equal(v_w, 0.5 * v_w_ + g_w)
  
      np.testing.assert_equal(w_g, w_0)
      np.testing.assert_equal(w_g, w_1)
      np.testing.assert_equal(b_g, b_0)
      np.testing.assert_equal(b_g, b_1)
  
      # Check params update step
      np.testing.assert_equal(w_0, w_g_ + v_w)
      np.testing.assert_equal(b_0, b_g_ + v_b)
  
imports: from multiprocessing import Process, Queue
  import numpy
  import os
  import shutil
  import tempfile
  import unittest
  import time
  from unittest.mock import Mock
  from hypothesis import assume, given, settings
  import hypothesis.strategies
  from caffe2.proto import caffe2_pb2
  from caffe2.python import brew, core, cnn, data_parallel_model, dyndep, model_helper, optimizer, rnn_cell, workspace
  from caffe2.python.test_util import TestCase
label_desc: Use image node_id 5 for calling the ParallelizeBMUFTest obj's underlying member method code with example usage: obj.test_parallelize_bmuf(cpu_device) and returns: None