Merge pull request #454 from probcomp/20220303-alexlew-homogeneous-mi…

…xture-grad Fix #445
probcomp · Mar 3, 2022 · 86f27bc · 86f27bc
2 parents 74b151e + ab0bfc5
commit 86f27bc
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/modeling_library/mixture.jl b/src/modeling_library/mixture.jl
@@ -114,7 +114,7 @@ function Gen.logpdf_grad(dist::HomogeneousMixture, x, weights, args...)
                     dims=dist.dims[i]+1)
                 grad_weights = reshape(
                     relative_weighted_densities,
-                    (1 for d in 1:dist.dims[i])..., length(dist.dims))
+                    (1 for d in 1:dist.dims[i])..., K)
             end
             push!(arg_grads, grads .* grad_weights)
         else