Merge pull request #747 from ufal/fix_stateful_context

Fix StatefulContext
ufal · Aug 26, 2018 · dd26a0f · dd26a0f
2 parents c59d1c9 + 7e2299b
commit dd26a0f
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/neuralmonkey/attention/stateful_context.py b/neuralmonkey/attention/stateful_context.py
@@ -63,7 +63,7 @@ def attention(self,
                                                            AttentionLoopState]:
         context = tf.reshape(self.attention_states,
                              [-1, self.context_vector_size])
-        weights = tf.ones(shape=[tf.shape(context)[0]])
+        weights = tf.ones(shape=[self.batch_size, 1])
 
         next_contexts = tf.concat(
             [loop_state.contexts, tf.expand_dims(context, 0)], 0)
@@ -77,8 +77,7 @@ def attention(self,
 
     def initial_loop_state(self) -> AttentionLoopState:
         return empty_attention_loop_state(
-            self.batch_size,
-            tf.shape(self.attention_states)[1],
+            self.batch_size, 1,
             self.context_vector_size)
 
     def finalize_loop(self, key: str,