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

Fix error with tf.function

Author: Rishit-dagli

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)

perceiver/version.py

@@ -1 +1 @@
-__version__ = "0.1.1"
+__version__ = "0.1.2"

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,