forgot the gpu only deco

pytorch · Nov 1, 2024 · 6f1d03f · 6f1d03f
1 parent 0702ef3
commit 6f1d03f
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/tests/torchtune/training/test_activation_offloading.py b/tests/torchtune/training/test_activation_offloading.py
@@ -134,6 +134,7 @@ def fwd(t):
     loss_c.backward()
 
 
+@gpu_test(gpu_count=1)
 def test_offloading_works_with_view_ac_cached_buffers() -> None:
     """
     Similar to test_offloading_works_with_view_outputs, but for when AC stashes