Update npu recompute unittests to cover offloading cases

PaddlePaddle · Jan 6, 2025 · e9c6439 · e9c6439
1 parent 43b1e12
commit e9c6439
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/backends/npu/tests/unittests/test_dygraph_recompute_for_eager.py b/backends/npu/tests/unittests/test_dygraph_recompute_for_eager.py
@@ -297,7 +297,7 @@ def test_recompute_kwargs(self):
         pos = paddle.randn(shape=[10, 10], dtype="float32")
         pos.stop_gradient = False
 
-        kwargs = {"pos": pos, "use_reentrant": True}
+        kwargs = {"pos": pos, "use_reentrant": True, "offload_indices": [0]}
         loss_ref, param_ref, grad_ref = run_model(
             recompute_block=[2], recompute_kwargs=kwargs
         )