Merge pull request #21 from Rishit-dagli/Rishit-dagli-patch-1

Fix error with tf.function
Rishit-dagli · Apr 26, 2021 · ca44960 · ca44960
2 parents 4d3b9b0 + c97c3b4
commit ca44960
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 3 deletions.
diff --git a/perceiver/perceiver.py b/perceiver/perceiver.py
@@ -76,6 +76,8 @@ def __init__(
             # self.existing_layers = get_latent_attn()(self.existing_layers)
             # self.existing_layers = get_latent_ff()(self.existing_layers)
 
+        self.existing_layers = tf.keras.Sequential(self.existing_layers)
+
         self.to_logits = tf.keras.Sequential(
             [
                 tf.keras.layers.LayerNormalization(axis=-1),
@@ -103,7 +105,7 @@ def call(self, data, mask=None):
 
         x = repeat(self.latents, "n d -> b n d", b=b)
 
-        x = tf.keras.Sequential(self.existing_layers)(x)
+        x = self.existing_layers(x)
 
         x = tf.math.reduce_mean(x, axis=-2)
         return self.to_logits(x)
diff --git a/perceiver/version.py b/perceiver/version.py
@@ -1 +1 @@
-__version__ = "0.1.1"
+__version__ = "0.1.2"
diff --git a/setup.py b/setup.py
@@ -7,7 +7,7 @@
 
 setup(
     name="perceiver",
-    version="0.1.1",
+    version="0.1.2",
     description="Implement of Perceiver, General Perception with Iterative Attention in TensorFlow",
     packages=["perceiver"],
     long_description=long_description,