code: def test_weight_scale(
      self, inputs, ITER, stepsize, upper_bound_iter, scale, gc, dc
  ):
      ITER = np.array([ITER], dtype=np.int64)
      op = core.CreateOperator(
          "WeightScale",
          ["w", "iter"],
          ["nw"],
          stepsize=stepsize,
          upper_bound_iter=upper_bound_iter,
          scale=scale,
      )
  
      def ref_weight_scale(w, iter, stepsize, upper_bound_iter, scale):
          iter = iter + 1
          return [
              w * scale
              if iter % stepsize == 0 and iter < upper_bound_iter
              else w
          ]
  
      input_device_options = {"iter": hu.cpu_do}
      self.assertReferenceChecks(
          gc,
          op,
          [inputs[0], ITER],
          functools.partial(
              ref_weight_scale,
              stepsize=stepsize,
              upper_bound_iter=upper_bound_iter,
              scale=scale,
          ),
          input_device_options=input_device_options,
      )
  
file_path: pytorch/caffe2/python/operator_test/weight_scale_test.py
imports: from caffe2.python import core
  import caffe2.python.hypothesis_test_util
  import functools
  from hypothesis import given
  import hypothesis.strategies
  import numpy
label_desc: Use image node_id 1 for calling the TestWeightScale obj's underlying member method code with example usage: obj.test_weight_scale(inputs, ITER, stepsize, upper_bound_iter, scale, gc, dc) and returns: list