[Sage/Uncertainty] Adds Gaussian process layer to CSF triggering head.

PiperOrigin-RevId: 332726662
google · Sep 20, 2020 · 02b60b8 · 02b60b8
1 parent ee83c00
commit 02b60b8
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/edward2/tensorflow/layers/random_feature.py b/edward2/tensorflow/layers/random_feature.py
@@ -185,6 +185,10 @@ def __init__(self,
   def build(self, input_shape):
     gp_feature_dim = input_shape[-1]
 
+    # Convert gp_feature_dim to int value for TF1 compatibility.
+    if isinstance(gp_feature_dim, tf.compat.v1.Dimension):
+      gp_feature_dim = gp_feature_dim.value
+
     # Posterior precision matrix for the GP' random feature coefficients.
     self.precision_matrix = (
         self.add_weight(