sum over the gamma distribution in the log_prior over last dim

gsbDBI · Dec 4, 2023 · f6e7570 · f6e7570
1 parent 09d3e78
commit f6e7570
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/bemb/model/bayesian_coefficient.py b/bemb/model/bayesian_coefficient.py
@@ -270,8 +270,8 @@ def log_prior(self,
             '''
             out = Gamma(concentration=concentration,
                         rate=rate).log_prob(sample)
-            # drop the last dim, take the first element over the last dim
-            out = out[:, :, 0]
+            # sum over the last dimension
+            out = torch.sum(out, dim=-1)
 
 
         # DEBUG_MARKER