diff options
| author |  Patrick Nguyen <drpng@google.com> | 2017-12-28 16:04:42 -0800 |
|---|---|---|
| committer |  TensorFlower Gardener <gardener@tensorflow.org> | 2017-12-28 16:08:58 -0800 |
| commit | 20765b3e1ae3b718699592c98aa9805cb874b6d1 (patch) | |
| tree | b429a74cd0046404644f34cc8fe6ff2cab78bb85 /tensorflow/contrib/memory_stats | |
| parent | 2e2715baa84720f786b38d1f9cb6887399020d6f (diff) | |
Merge changes from github.
PiperOrigin-RevId: 180301735
Diffstat (limited to 'tensorflow/contrib/memory_stats')
| -rw-r--r-- | tensorflow/contrib/memory_stats/__init__.py | 4 | ||||
| -rw-r--r-- | tensorflow/contrib/memory_stats/python/kernel_tests/memory_stats_ops_test.py | 5 |
2 files changed, 7 insertions, 2 deletions
diff --git a/tensorflow/contrib/memory_stats/__init__.py b/tensorflow/contrib/memory_stats/__init__.py index a32302c854..2ce849ca66 100644 --- a/tensorflow/contrib/memory_stats/__init__.py +++ b/tensorflow/contrib/memory_stats/__init__.py @@ -19,6 +19,10 @@ @@MaxBytesInUse """ +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + from tensorflow.contrib.memory_stats.python.ops.memory_stats_ops import BytesInUse from tensorflow.contrib.memory_stats.python.ops.memory_stats_ops import BytesLimit from tensorflow.contrib.memory_stats.python.ops.memory_stats_ops import MaxBytesInUse diff --git a/tensorflow/contrib/memory_stats/python/kernel_tests/memory_stats_ops_test.py b/tensorflow/contrib/memory_stats/python/kernel_tests/memory_stats_ops_test.py index 5e52ef3647..02c2ac06fb 100644 --- a/tensorflow/contrib/memory_stats/python/kernel_tests/memory_stats_ops_test.py +++ b/tensorflow/contrib/memory_stats/python/kernel_tests/memory_stats_ops_test.py @@ -76,9 +76,10 @@ class MemoryStatsOpsTest(test_util.TensorFlowTestCase): with ops.control_dependencies([a]): bytes_in_use_op = memory_stats_ops.BytesInUse() with ops.control_dependencies([bytes_in_use_op]): - b = math_ops.add(a, a) + b = random_ops.random_uniform(matrix_shape, dtype=dtype) + c = math_ops.matmul(a, b) - _, bytes_in_use, max_bytes_in_use = sess.run([b, bytes_in_use_op, + _, bytes_in_use, max_bytes_in_use = sess.run([c, bytes_in_use_op, max_bytes_in_use_op]) # intermediate result allocates 1 matrix, max usage is at least 2 |