fix calculating E[x2] in FullSupportBarDistribution

PriorLabs · Feb 5, 2025 · 4c79b45 · 4c79b45
1 parent f1af310
commit 4c79b45
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/tabpfn/model/bar_distribution.py b/src/tabpfn/model/bar_distribution.py
@@ -620,7 +620,7 @@ def mean_of_square(self, logits: torch.Tensor) -> torch.Tensor:
         )
         bucket_mean_of_square[-1] = (
             side_normals[1].variance
-            + (side_normals[1].variance + self.borders[-2]).square()
+            + (side_normals[1].mean + self.borders[-2]).square()
         )
         p = torch.softmax(logits, -1)
         return p @ bucket_mean_of_square