text: 
  
  
  
  
  from hypothesis import given
  import numpy as np
  import hypothesis.strategies as st
  
  from caffe2.python import core, workspace
  import caffe2.python.hypothesis_test_util as hu
  
  
  @st.composite
  def _dev_options(draw):
      op_dev = draw(st.sampled_from(hu.device_options))
      if op_dev == hu.cpu_do:
          # the CPU op can only handle CPU tensor
          input_blob_dev = hu.cpu_do
      else:
          input_blob_dev = draw(st.sampled_from(hu.device_options))
  
      return op_dev, input_blob_dev
  
  
  class TestEnsureCPUOutputOp(hu.HypothesisTestCase):
  
      @given(
          input=hu.tensor(dtype=np.float32),
          dev_options=_dev_options()
      )
      def test_ensure_cpu_output(self, input, dev_options):
          op_dev, input_blob_dev = dev_options
          net = core.Net('test_net')
          data = net.GivenTensorFill(
              [],
              ["data"],
              values=input,
              shape=input.shape,
              device_option=input_blob_dev
          )
  
          data_cpu = net.EnsureCPUOutput(
              [data],
              ["data_cpu"],
              device_option=op_dev
          )
          workspace.RunNetOnce(net)
  
          data_cpu_value = workspace.FetchBlob(data_cpu)
          np.testing.assert_allclose(input, data_cpu_value)
  