Apply temporal merge to the time dimensions, not the cell/batch dimen…

…sion
vanvalenlab · Jun 5, 2024 · 09ba804 · 09ba804
1 parent e6d7af8
commit 09ba804
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/deepcell/layers/temporal.py b/deepcell/layers/temporal.py
@@ -97,7 +97,7 @@ def __init__(self, encoder_dim=64, **kwargs):
     def call(self, inputs):
         input_shape = tf.shape(inputs)
         # reshape away the temporal axis
-        x = tf.reshape(inputs, [-1, input_shape[2], self.encoder_dim])
+        x = tf.reshape(inputs, [-1, input_shape[1], self.encoder_dim])
         x = self.lstm(x)
         output_shape = [-1, input_shape[1], input_shape[2], self.encoder_dim]
         x = tf.reshape(x, output_shape)