Replace references to deprecated jax array attributes device_buffer a…

…nd device_buffers PiperOrigin-RevId: 588553845
google-research · Dec 6, 2023 · 62c75bc · 62c75bc
1 parent 1bfd2f1
commit 62c75bc
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/t5x/trainer_test.py b/t5x/trainer_test.py
@@ -175,7 +175,7 @@ def test_timer_blocking_on_donated_buffer(self):
     mm._duration_timer._start_future.result()
 
     # Deleted/donated.
-    x.device_buffer.delete()
+    x.addressable_data(0).delete()
     mm.start_duration_timer(block_on=x)
     mm._duration_timer._start_future.result()