[FIX] Disable argument donation (again)

neel04 · Dec 19, 2024 · b90f471 · b90f471
1 parent 65aee4a
commit b90f471
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/Dockerfile b/Dockerfile
@@ -26,3 +26,4 @@ WORKDIR /ReAct_Jax
 
 # Set the entry point to bash 
 ENTRYPOINT ["/bin/bash"]
+
diff --git a/ReAct/utils/trainer.py b/ReAct/utils/trainer.py
@@ -71,7 +71,7 @@ def _compute_softmax_cross_entropy_loss(pred_y: Array, y_one_hot: Array) -> Arra
 
     return loss.mean()
 
-@eqx.filter_jit(donate='all-except-first')
+@eqx.filter_jit
 def make_step(
     keys: PRNGKeyArray,
     model: Union[React, GPT],
Original file line number	Diff line number	Diff line change
Expand Up		@@ -26,3 +26,4 @@ WORKDIR /ReAct_Jax

		# Set the entry point to bash
		ENTRYPOINT ["/bin/bash"]