class_docstr: Create and run the model. We try with both storing indices for gather
  on CPU and on GPU
code: def run_model(self, V, gpu_devices, cpu_indices):
      def input_builder_fun(model):
          return None
  
      def model_build_fun(model, loss_scale):
          if cpu_indices:
              with core.DeviceScope(core.DeviceOption(caffe2_pb2.CPU)):
                  gathered_cpu = model.net.Gather(
                      [self.vecs, "indices"], "gathered_cpu"
                  )
  
              gathered = model.CopyCPUToGPU(gathered_cpu, "gathered")
          else:
              gpu_vecs = model.param_init_net.CopyCPUToGPU(
                  self.vecs,
                  "gpuvecs",
              )
              model.params.append(gpu_vecs)
              gathered = model.net.Gather(
                  [gpu_vecs, "indices"], "gathered"
              )
          flattened = model.Flatten(gathered, "flattened")
          fc = model.FC(
              flattened,
              "fc",
              16 * 16,
              1,
              ("ConstantFill", {}),
              ("ConstantFill", {}),
          )
          fc_fl = model.FlattenToVec(fc, "fc_fl")
          sigm = model.Sigmoid(fc_fl, "sigm")
          sq = model.SquaredL2Distance([sigm, "label"], "sq")
          loss = model.AveragedLoss(sq, "loss")
          loss = model.Scale(loss, scale=loss_scale)
          return [loss]
  
      def param_update_fun(model):
          ONE = model.param_init_net.ConstantFill(
              [],
              "ONE",
              shape=[1],
              value=1.0,
          )
          LR = model.CopyCPUToGPU(self.LR, "LR")
          for param in model.GetParams():
              param_grad = model.param_to_grad[param]
              if not isinstance(param_grad, core.GradientSlice):
                  model.WeightedSum([param, ONE, param_grad, LR], param)
              else:
                  param_momentum = model.param_init_net.ConstantFill(
                      [param],
                      param + "_momentum",
                      value=0.0,
                  )
                  model.net.SparseMomentumSGDUpdate(
                      [
                          param_grad.values,
                          param_momentum,
                          LR,
                          param,
                          param_grad.indices,
                      ],
                      [param_grad.values, param_momentum, param],
                      momentum=0.1,
                      nesterov=0,
                  )
  
      workspace.ResetWorkspace()
      model = cnn.CNNModelHelper(
          order="NHWC",
          name="sparse_test{}".format(gpu_devices),
      )
  
      with core.NameScope("cpu"):
          with core.DeviceScope(core.DeviceOption(caffe2_pb2.CPU)):
              self.ITER = model.Iter("ITER")
              self.LR = model.net.LearningRate(
                  [self.ITER],
                  "LR",
                  base_lr=(-0.1),
                  policy="fixed",
              )
              self.vecs = model.param_init_net.UniformFill(
                  [], "vecs", shape=[V, 16]
              )
              if cpu_indices:
                  model.params.append(self.vecs)
              self.ONE_CPU = model.param_init_net.ConstantFill(
                  [],
                  "ONE_CPU",
                  shape=[1],
                  value=1.0,
              )
  
      data_parallel_model.Parallelize_GPU(
          model,
          input_builder_fun=input_builder_fun,
          forward_pass_builder_fun=model_build_fun,
          param_update_builder_fun=param_update_fun,
          devices=gpu_devices,
      )
  
      # Update the vecs
      if cpu_indices:
          with core.NameScope("cpu"):
              with core.DeviceScope(core.DeviceOption(caffe2_pb2.CPU)):
                  for param in model.GetParams():
                      param_grad = model.param_to_grad[param]
                      model.ScatterWeightedSum(
                          [
                              param,
                              self.ONE_CPU,
                              param_grad.indices,
                              param_grad.values,
                              self.LR,
                          ],
                          self.vecs,
                      )
      else:
          with core.DeviceScope(
              core.DeviceOption(workspace.GpuDeviceType, 0)
          ):
              model.CopyGPUToCPU("gpu_0/gpuvecs", self.vecs)
  
      np.random.seed(2603)
  
      # Each run has same input, independent of number of gpus
      batch_size = 64
      for i in range(0, 10):
          full_indices = np.random.permutation(V)[
              : batch_size * 16
          ].reshape(batch_size, 16)
          full_labels = full_indices[:, 0] % 2
          batch_per_device = batch_size // len(gpu_devices)
  
          for j, g in enumerate(gpu_devices):
              st = j * batch_per_device
              en = st + batch_per_device
              indices = full_indices[st:en, :].astype(np.int32)
              labels = full_labels[st:en].astype(np.float32)
  
              device_for_indices = core.DeviceOption(caffe2_pb2.CPU)
              if not cpu_indices:
                  device_for_indices = core.DeviceOption(
                      workspace.GpuDeviceType, g
                  )
  
              with core.DeviceScope(device_for_indices):
                  workspace.FeedBlob(
                      "gpu_{}/indices".format(g), indices
                  )
  
              with core.DeviceScope(
                  core.DeviceOption(workspace.GpuDeviceType, g)
              ):
                  workspace.FeedBlob("gpu_{}/label".format(g), labels)
  
          if i == 0:
              workspace.RunNetOnce(model.param_init_net)
              # Force vecs to be same on all runs
              orig_vecs = np.random.rand(V, 16).astype(np.float32)
              workspace.FeedBlob(self.vecs, orig_vecs)
              if not cpu_indices:
                  for g in gpu_devices:
                      workspace.FeedBlob(
                          "gpu_{}/gpuvecs".format(g),
                          orig_vecs,
                          device_option=core.DeviceOption(
                              workspace.GpuDeviceType, g
                          ),
                      )
              workspace.CreateNet(model.net)
  
          workspace.RunNet(model.net.Proto().name)
          if len(gpu_devices) == 2:
              if not cpu_indices:
                  idx = workspace.FetchBlob("gpu_0/indices")
                  idx = list(idx.flatten())
                  n = len(idx)
                  nu = len(set(idx))
                  assert n == nu, "We cannot have duplicate indices"
  
      # Sanity check to see the vecs were updated
      self.assertFalse(
          np.allclose(workspace.FetchBlob(self.vecs), orig_vecs)
      )
      return [
          workspace.FetchBlob(
              self.vecs if cpu_indices else "gpu_0/gpuvecs"
          ),
          workspace.FetchBlob("gpu_0/fc_w"),
      ]
  
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 1 for calling the SparseDataParallelModelTest obj's underlying member method code with example usage: obj.run_model(V, gpu_devices, cpu_indices) and returns: list, None, list