diff --git a/phygnn/layers/custom_layers.py b/phygnn/layers/custom_layers.py
index 7f1af94..f0e0b13 100644
--- a/phygnn/layers/custom_layers.py
+++ b/phygnn/layers/custom_layers.py
@@ -978,6 +978,17 @@ def __init__(self, name=None, adder=0, inverse=False):
         super().__init__(name=name)
         self.adder = adder
         self.inverse = inverse
+        self.rank = None
+
+    def build(self, input_shape):
+        """Custom implementation of the tf layer build method.
+
+        Parameters
+        ----------
+        input_shape : tuple
+            Shape tuple of the input
+        """
+        self.rank = len(input_shape)
 
     def call(self, x):
         """Operates on x with (inverse) log transform