Unverified Commit 8d82db43 authored by Edward Z. Yang's avatar Edward Z. Yang Committed by GitHub
Browse files

Future proof storage size test (#735)

See https://github.com/pytorch/pytorch/pull/59671/

Signed-off-by: default avatarEdward Z. Yang <ezyang@fb.com>
parent d442ad18
......@@ -31,7 +31,12 @@ def test_memory_leak():
bucket.add_grad(param)
bucket.shrink()
assert len(bucket.buffer.storage()) == 6
storage = bucket.buffer.storage()
# See https://github.com/pytorch/pytorch/pull/59671/
if hasattr(storage, "nbytes"):
assert storage.nbytes() == 6 * bucket.buffer.element_size()
else:
assert len(storage) == 6
def test_max_size():
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment