diff options
| author |  A. Unique TensorFlower <gardener@tensorflow.org> | 2017-06-09 11:26:38 -0700 |
|---|---|---|
| committer | TensorFlower Gardener <gardener@tensorflow.org> | 2017-06-09 11:29:40 -0700 |
| commit | e64b6ee7e5ed5571b2d821d3a85e8519ef807c49 (patch) | |
| tree | 55b5e5ccb4dad0fb3bef21eb15c2c59a67c92128 /tensorflow/contrib/memory_stats | |
| parent | c60899ae706357383826ea05cf380462eb3fe2ad (diff) | |
Add static shape inference to BytesLimit and MaxBytesInUse.
PiperOrigin-RevId: 158540875
Diffstat (limited to 'tensorflow/contrib/memory_stats')
| -rw-r--r-- | tensorflow/contrib/memory_stats/ops/memory_stats_ops.cc | 11 |
1 files changed, 9 insertions, 2 deletions
diff --git a/tensorflow/contrib/memory_stats/ops/memory_stats_ops.cc b/tensorflow/contrib/memory_stats/ops/memory_stats_ops.cc index d5c6c4ea87..08859c8613 100644 --- a/tensorflow/contrib/memory_stats/ops/memory_stats_ops.cc +++ b/tensorflow/contrib/memory_stats/ops/memory_stats_ops.cc @@ -12,11 +12,18 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. ==============================================================================*/ +#include "tensorflow/core/framework/common_shape_fns.h" #include "tensorflow/core/framework/op.h" namespace tensorflow { -REGISTER_OP("BytesLimit").Output("out: int64").SetIsStateful(); -REGISTER_OP("MaxBytesInUse").Output("out: int64").SetIsStateful(); +REGISTER_OP("BytesLimit") + .Output("out: int64") + .SetIsStateful() + .SetShapeFn(shape_inference::ScalarShape); +REGISTER_OP("MaxBytesInUse") + .Output("out: int64") + .SetIsStateful() + .SetShapeFn(shape_inference::ScalarShape); } // namespace tensorflow |