Add logging on the allocations

apache · May 9, 2020 · 4e61242 · 4e61242
1 parent ab7cd6d
commit 4e61242
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/tests/python/unittest/test_memory_opt.py b/tests/python/unittest/test_memory_opt.py
@@ -182,10 +182,12 @@ def test_resnet152():
                                              data=(32, 3, 224, 224))
     os.environ["MXNET_MEMORY_OPT"] = '0'
     no_opt_exec = resnet_152.simple_bind(mx.cpu(), 'write', data=(32, 3, 224, 224))
-    assert grep_exec_memory_consumption(memory_opt_exec) /  \
-           grep_exec_memory_consumption(no_opt_exec) < 0.6, \
+    memory_opt_alloc = grep_exec_memory_consumption(memory_opt_exec)
+    no_opt_alloc = grep_exec_memory_consumption(no_opt_exec)
+    assert memory_opt_alloc / no_opt_alloc < 0.6, \
            "The ratio between the memory consumption with the memory optimizations " \
-           "enabled and disabled is expected to be smaller than 0.6"
+           "enabled and disabled (%d vs. %d MB) is expected to be smaller than 0.6"  \
+           % (memory_opt_alloc, no_opt_alloc)
 
 
 if __name__ == "__main__":