refactor rng_fn method (#212)

aloctavodia · web-flow · commit cd5dfbe4e09e · 2024-12-29T08:11:49.000-03:00
diff --git a/pymc_bart/bart.py b/pymc_bart/bart.py
@@ -55,12 +55,12 @@ def rng_fn(  # pylint: disable=W0237
         if not size:
             size = None
 
-        if isinstance(cls.Y, (TensorSharedVariable, TensorVariable)):
-            Y = cls.Y.eval()
-        else:
-            Y = cls.Y
-
         if not cls.all_trees:
+            if isinstance(cls.Y, (TensorSharedVariable, TensorVariable)):
+                Y = cls.Y.eval()
+            else:
+                Y = cls.Y
+
             if size is not None:
                 return np.full((size[0], Y.shape[0]), Y.mean())
             else:
