Update Docker to accomodate jax rollback, fixing #21116

neel04 · May 10, 2024 · c7e4033 · c7e4033
1 parent a6d53e8
commit c7e4033
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/Dockerfile b/Dockerfile
@@ -24,5 +24,5 @@ RUN pip3 install git+https://github.com/deepmind/jmp
 
 WORKDIR /ReAct_Jax
 
-# Set the entry point to bash
+# Set the entry point to bash 
 ENTRYPOINT ["/bin/bash"]
diff --git a/ReAct/model/react.py b/ReAct/model/react.py
@@ -79,7 +79,7 @@ def f(input_tup: Tuple[Array, int], _dynamic_bl: PyTree) -> Tuple[Tuple[Array, i
 
         #ctx_state *= jax.nn.sigmoid(self.forget_gate(hist_lerp, enable_dropout, key))
         #ctx_state += self.ctx_gate(hist_lerp, enable_dropout, key)
-        ctx_state = self.ctx_gate(hist_lerp)
+        ctx_state += self.ctx_gate(hist_lerp)
 
         return out[0], ctx_state