[TRAX] Fix the jnp.pad issue with mode.

PiperOrigin-RevId: 369691659
google · Apr 21, 2021 · df96ed6 · df96ed6
1 parent 65378ce
commit df96ed6
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/trax/layers/research/sparsity.py b/trax/layers/research/sparsity.py
@@ -243,7 +243,7 @@ def forward(self, x):
     else:
       pad_widths = [[0, 0] for _ in range(len(x.shape))]
       pad_widths[1][0] = self._n_items_to_remember
-      x = jnp.pad(x, pad_width=pad_widths)
+      x = jnp.pad(x, pad_width=pad_widths, mode='constant')
     return x
 
   def init_weights_and_state(self, input_signature):