| 295 | |
| 296 | template <typename TKey, typename TValue, typename TGlobalStep> |
| 297 | void DumpEv( |
| 298 | OpKernelContext* context, |
| 299 | EmbeddingVar<TKey, TValue>* variable, |
| 300 | const string& tensor_name, BundleWriter& writer) { |
| 301 | const Tensor& global_step = context->input(5); |
| 302 | TGlobalStep global_step_scalar = global_step.scalar<TGlobalStep>()(); |
| 303 | core::ScopedUnref s(variable); |
| 304 | embedding::ShrinkArgs shrink_args; |
| 305 | shrink_args.global_step = global_step_scalar; |
| 306 | const Tensor& prefix = context->input(0); |
| 307 | const string& prefix_string = prefix.scalar<tstring>()(); |
| 308 | OP_REQUIRES_OK(context, variable->Save(tensor_name, |
| 309 | prefix_string, &writer, shrink_args)); |
| 310 | } |
| 311 | |
| 312 | void Compute(OpKernelContext* context) override { |
| 313 | const Tensor& prefix = context->input(0); |