code: def test_forward_optim_tree_harder(
      self, input_dim, output_dim, batch_size
  ):
      m = model_helper.ModelHelper()
      m.net.Proto().type = "dag"
      m.net.Proto().num_workers = 4
      m.net.AddExternalInput("label")
      m.net.AddExternalInput("data")
  
      with core.NameScope("name_x"):
          fc1 = brew.fc(
              m, "data", "fc1", dim_in=input_dim, dim_out=output_dim
          )
          fc2 = brew.fc(
              m, fc1, "fc2", dim_in=output_dim, dim_out=output_dim
          )
  
          fc3 = brew.fc(
              m, fc2, "fc3", dim_in=output_dim, dim_out=output_dim
          )
          fc4 = brew.fc(
              m, fc3, "fc4", dim_in=output_dim, dim_out=output_dim
          )
          fc5 = brew.fc(
              m, fc4, "fc5", dim_in=output_dim, dim_out=output_dim
          )
  
          # Branch
          fc3b = brew.fc(
              m, fc2, "fc3b", dim_in=output_dim, dim_out=output_dim
          )
          fc4b = brew.fc(
              m, fc3b, "fc4b", dim_in=output_dim, dim_out=output_dim
          )
          fc5b = brew.fc(
              m, fc4b, "fc5b", dim_in=output_dim, dim_out=output_dim
          )
  
          fc5sum = brew.sum(m, [fc5, fc5b], "fc5sum")
          fc5sum.Relu([], "relu1").Softmax(
              [], "pred1"
          ).LabelCrossEntropy(["label"], ["xent1"]).AveragedLoss(
              [], "loss1"
          )
          fc6 = brew.fc(
              m, fc5, "fc6", dim_in=output_dim, dim_out=output_dim
          )
          fc6.Relu([], fc6).Softmax([], "pred2").LabelCrossEntropy(
              ["label"], ["xent2"]
          ).AveragedLoss([], "loss2")
  
      blobs_before = count_blobs(m.net.Proto())
      optim_proto = memonger.optimize_inference_for_dag(
          m.net, ["name_x/data"], "name_x/"
      )
  
      blobs_after = count_blobs(optim_proto)
  
      # Extra test with when one of the parameters is also an input.
      # This caused a bug before.
      optim_proto_extra_input = memonger.optimize_inference_for_dag(
          m.net, ["name_x/data", "name_x/fc1_w"], "name_x/"
      )
      blobs_after_extra_input = count_blobs(optim_proto_extra_input)
      self.assertEqual(blobs_after, blobs_after_extra_input)
      ###
  
      print(str(optim_proto))
      self.assertLess(blobs_after, blobs_before)
  
      # Test networks produce exactly same results
      data = np.random.randn(batch_size, input_dim).astype(np.float32)
      label = np.random.randint(
          low=0, high=output_dim, size=(batch_size,)
      ).astype(np.int32)
      workspace.RunNetOnce(m.param_init_net)
      workspace.FeedBlob("name_x/data", data)
      workspace.FeedBlob("name_x/label", label)
      workspace.RunNetOnce(m.net)
      loss1 = workspace.FetchBlob("name_x/loss1")
      loss2 = workspace.FetchBlob("name_x/loss2")
      workspace.RunNetOnce(optim_proto)
      optimized_loss1 = workspace.FetchBlob("name_x/loss1")
      optimized_loss2 = workspace.FetchBlob("name_x/loss2")
      np.testing.assert_almost_equal(loss1, optimized_loss1)
      np.testing.assert_almost_equal(loss2, optimized_loss2)
  
imports: import numpy
  from caffe2.python import workspace, memonger, core, model_helper, brew
  from caffe2.proto import caffe2_pb2
  import caffe2.python.hypothesis_test_util
  import hypothesis.strategies
  from hypothesis import given, settings
  import unittest
label_desc: Use image node_id 8 for calling the MemongerTest obj's underlying member method code with example usage: obj.test_forward_optim_tree_harder(input_dim, output_dim, batch_size) without return types