Fix bytes_limit key of mem_stats (#243)

ROCm · Feb 26, 2025 · c5ee54f · c5ee54f
1 parent bbaaec0
commit c5ee54f
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/multi_device_test.py b/tests/multi_device_test.py
@@ -312,7 +312,7 @@ def test_lax_full_like_efficient(self):
       self.skipTest('Only can run test on device with mem_stats')
     mesh = Mesh(devices, axis_names=("i"))
     sharding = NamedSharding(mesh, P('i'))
-    available_memory = mem_stats['bytes_reservable_limit']
+    available_memory = mem_stats['bytes_limit']
     array_size = available_memory // (6 * len(devices)) * len(devices)
     # Set up tracemalloc to track memory usage.
     tm.start()