More minor fixes

hmsc-r · Dec 4, 2023 · c92c475 · c92c475
1 parent 7724142
commit c92c475
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/hmsc/updaters/updateZ.py b/hmsc/updaters/updateZ.py
@@ -43,7 +43,7 @@ def updateZ(params, data, rLHyperparams, *,
     Pi = data["Pi"]
     distr = data["distr"]
 
-    if X.ndim == 2:
+    if len(X.shape.as_list()) == 2: #tf.rank(X) X.ndim == 2:
         LFix = tf.matmul(X, Beta)
     else:
         LFix = tf.einsum("jik,kj->ij", X, Beta)