refactor rng_fn method (#212)

pymc-devs · Dec 29, 2024 · cd5dfbe · cd5dfbe
1 parent 064457e
commit cd5dfbe
Showing 1 changed file with 5 additions and 5 deletions.
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: