From 607816029977fd29808786b4305551c992733842 Mon Sep 17 00:00:00 2001 From: Eugene Brevdo Date: Tue, 5 Sep 2017 16:37:32 -0700 Subject: [PATCH] Extended ScratchSpace to expose its underlying scratch tensor object. PiperOrigin-RevId: 167649551 --- tensorflow/core/kernels/cuda_solvers.h | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tensorflow/core/kernels/cuda_solvers.h b/tensorflow/core/kernels/cuda_solvers.h index ac6119d8a21..0fd6450f982 100644 --- a/tensorflow/core/kernels/cuda_solvers.h +++ b/tensorflow/core/kernels/cuda_solvers.h @@ -313,6 +313,9 @@ class ScratchSpace { int64 size() const { return scratch_tensor_.NumElements(); } const string& debug_info() const { return debug_info_; } + Tensor& tensor() { return scratch_tensor_; } + const Tensor& tensor() const { return scratch_tensor_; } + // Returns true if this ScratchSpace is in host memory. bool on_host() const { return on_host_; }