use parent when fill! a triangular matrix

TuringLang · Jul 25, 2023 · 8133567 · 8133567
1 parent 9cd5907
commit 8133567
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/utils.jl b/src/utils.jl
@@ -30,7 +30,7 @@ function triu_mask(X::AbstractMatrix, k::Int)
 
  # Using `similar` allows us to respect device of array, etc., e.g. `CuArray`.
  m = similar(X, Bool)
- return triu!(fill!(m, true), k)
+ return triu!(fill!(parent(m), true), k)
 end
 
 ChainRulesCore.@non_differentiable triu_mask(X::AbstractMatrix, k::Int)