back to tf.tensorshape

innat · Mar 21, 2024 · 86020bd · 86020bd
1 parent 7e61d14
commit 86020bd
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/videoswin/blocks/basic.py b/videoswin/blocks/basic.py
@@ -118,14 +118,14 @@ def compute_output_shape(self, input_shape):
         if self.downsample is not None:
             # TODO: remove tensorflow dependencies.
             # GitHub issue: https://github.com/keras-team/keras/issues/19259 # noqa: E501
-            output_shape = keras.KerasTensor(
+            output_shape = tf.TensorShape(
                 [
                     input_shape[0],
                     self.depth_pad,
                     self.height_pad // 2,
                     self.width_pad // 2,
                     2 * self.input_dim,
-                ], dtype=self.compute_dtype
+                ]
             )
             return output_shape