diff --git a/python/rmm/tests/test_rmm.py b/python/rmm/tests/test_rmm.py index b2ce8fe6d..b5dc81c1f 100644 --- a/python/rmm/tests/test_rmm.py +++ b/python/rmm/tests/test_rmm.py @@ -646,19 +646,19 @@ def test_statistics_resource_adaptor(stats_mr): buffers.append(rmm.DeviceBuffer(size=1000)) assert mr2.allocation_counts == { - "current_bytes": 2000, + "current_bytes": 2016, "current_count": 2, - "peak_bytes": 2000, + "peak_bytes": 2016, "peak_count": 2, - "total_bytes": 2000, + "total_bytes": 2016, "total_count": 2, } assert stats_mr.allocation_counts == { - "current_bytes": 7000, + "current_bytes": 7056, "current_count": 7, - "peak_bytes": 10000, + "peak_bytes": 10080, "peak_count": 10, - "total_bytes": 12000, + "total_bytes": 12096, "total_count": 12, } @@ -668,17 +668,17 @@ def test_statistics_resource_adaptor(stats_mr): assert mr2.allocation_counts == { "current_bytes": 0, "current_count": 0, - "peak_bytes": 2000, + "peak_bytes": 2016, "peak_count": 2, - "total_bytes": 2000, + "total_bytes": 2016, "total_count": 2, } assert stats_mr.allocation_counts == { "current_bytes": 0, "current_count": 0, - "peak_bytes": 10000, + "peak_bytes": 10080, "peak_count": 10, - "total_bytes": 12000, + "total_bytes": 12096, "total_count": 12, } gc.collect() @@ -696,7 +696,7 @@ def test_tracking_resource_adaptor(): for i in range(9, 0, -2): del buffers[i] - assert mr.get_allocated_bytes() == 5000 + assert mr.get_allocated_bytes() == 5040 # Push a new Tracking adaptor mr2 = rmm.mr.TrackingResourceAdaptor(mr, capture_stacks=True) @@ -705,8 +705,8 @@ def test_tracking_resource_adaptor(): for _ in range(2): buffers.append(rmm.DeviceBuffer(size=1000)) - assert mr2.get_allocated_bytes() == 2000 - assert mr.get_allocated_bytes() == 7000 + assert mr2.get_allocated_bytes() == 2016 + assert mr.get_allocated_bytes() == 7056 # Ensure we get back a non-empty string for the allocations assert len(mr.get_outstanding_allocations_str()) > 0