fix remat checkpoint of input

apple · Nov 21, 2024 · 408033d · 408033d
1 parent 45d769d
commit 408033d
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/axlearn/common/attention.py b/axlearn/common/attention.py
@@ -3159,7 +3159,7 @@ def _forward_for_mode(
         Raises:
             ValueError: If `mode` is unsupported.
         """
-        self._remat_name(data, "input")
+        data = self._remat_name(data, "input")
         self.vlog(3, "transformer.input=%s", data.sum())
         self_attention_return_aux = set()
         cross_attention_return_aux = set()