Update tacotron.py

tacotron.py

@@ -71,7 +71,7 @@ class BiGRU(pax.Module):
                 return jnp.where(reset_mask, x0, xt)
 
             state, _ = self.rnn_bwd(prev, x)
-            state = jax.tree_map(reset_state, x_bwd_states0, state)
+            state = jax.tree_util.tree_map(reset_state, x_bwd_states0, state)
             return state, state.hidden
 
         x_bwd_states, x_bwd = pax.scan(