| 510 | } |
| 511 | |
| 512 | void PrefetchDatasetOp::MakeDataset(OpKernelContext* ctx, DatasetBase* input, |
| 513 | DatasetBase** output) { |
| 514 | int64 buffer_size = 0; |
| 515 | OP_REQUIRES_OK(ctx, |
| 516 | ParseScalarArgument<int64>(ctx, kBufferSize, &buffer_size)); |
| 517 | OP_REQUIRES(ctx, buffer_size >= 0 || buffer_size == model::kAutotune, |
| 518 | errors::InvalidArgument("buffer_size must be >= 0 or set " |
| 519 | "buffer_size to be ", |
| 520 | model::kAutotune, " for auto-tuning")); |
| 521 | |
| 522 | if (buffer_size == model::kAutotune) { |
| 523 | metrics::RecordTFDataAutotune(kDatasetType); |
| 524 | } |
| 525 | |
| 526 | *output = |
| 527 | new Dataset(ctx, input, buffer_size, slack_period_, legacy_autotune_); |
| 528 | } |
| 529 | |
| 530 | namespace { |
| 531 | REGISTER_KERNEL_BUILDER(Name("PrefetchDataset").Device(DEVICE_CPU).Priority(2), |
nothing calls this directly
no test coverage detected