Fixed a bug in integral computation batchification for DiracMixture

Xeadriel · Nov 30, 2022 · 25cb72b · 25cb72b
1 parent 0b6d9bb
commit 25cb72b
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/firelang/measure/dirac.py b/firelang/measure/dirac.py
@@ -71,7 +71,7 @@ def integral(
             func = func.view(*func.shape[:-1], func_size, 1, 1)
 
             size = func_size * self.k
-            nrow_per_batch = (batch_size - size + 1) // size
+            nrow_per_batch = (batch_size + size - 1) // size
             fx = []
             for i in range(0, measure_size, nrow_per_batch):
                 _x = x[..., i : i + nrow_per_batch, :, :]