diff --git a/t5x/trainer_test.py b/t5x/trainer_test.py
index 0cd576b0b..f1a048e8d 100644
--- 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()