ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_);
  TF_GraphCopyFunction(host_graph_, func_, nullptr, s_);
  ASSERT_EQ(TF_OK, TF_GetCode(s_)) << TF_Message(s_);

  // Delete the gradient func.
  // It is safe to delete after adding a copy to host graph.
  TF_DeleteFunction(grad_func);

  // Check that GraphDef did not change
  GraphDef gdef2;
  GetGraphDef(host_graph_, &gdef2);

  for (int64_t i = 0; i < num_dims; ++i) {
    num_elements *= dims[i];
  }
  TF_Tensor* t = TF_AllocateTensor(TF_STRING, dims, num_dims,
                                   sizeof(TF_TString) * num_elements);
  delete[] dims;

  TF_TString* data = static_cast<TF_TString*>(TF_TensorData(t));
  for (int i = 0; i < batch_size; ++i) {
    TF_TString_Init(&data[i]);