Searched refs:batch_sizes_tensor (Results 1 – 1 of 1) sorted by relevance
296 const Tensor* batch_sizes_tensor; in MakeDataset() local297 OP_REQUIRES_OK(ctx, ctx->input("batch_sizes", &batch_sizes_tensor)); in MakeDataset()299 ctx, batch_sizes_tensor->dims() <= 1, in MakeDataset()303 batch_sizes.reserve(batch_sizes_tensor->NumElements()); in MakeDataset()304 for (int i = 0; i < batch_sizes_tensor->NumElements(); ++i) { in MakeDataset()305 batch_sizes.push_back(batch_sizes_tensor->flat<int64_t>()(i)); in MakeDataset()