Method Compute

tensorflow/core/kernels/summary_kernels.cc:329–343

  explicit WriteGraphSummaryOp(OpKernelConstruction* ctx) : OpKernel(ctx) {}

  void Compute(OpKernelContext* ctx) override {
    core::RefCountPtr<SummaryWriterInterface> s;
    OP_REQUIRES_OK(ctx, LookupResource(ctx, HandleFromInput(ctx, 0), &s));
    const Tensor* t;
    OP_REQUIRES_OK(ctx, ctx->input("step", &t));
    const int64 step = t->scalar<int64>()();
    OP_REQUIRES_OK(ctx, ctx->input("tensor", &t));
    std::unique_ptr<GraphDef> graph{new GraphDef};
    if (!ParseProtoUnlimited(graph.get(), t->scalar<tstring>()())) {
      ctx->CtxFailureWithWarning(
          errors::DataLoss("Bad tf.GraphDef binary proto tensor string"));
      return;
    }
    OP_REQUIRES_OK(ctx, s->WriteGraph(step, std::move(graph)));
  }
};
REGISTER_KERNEL_BUILDER(Name("WriteGraphSummary").Device(DEVICE_CPU),
                        WriteGraphSummaryOp);