reshape to unsqueeze

MurrellGroup · Feb 8, 2024 · 2b3fed5 · 2b3fed5
1 parent e859a6a
commit 2b3fed5
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/layers.jl b/src/layers.jl
@@ -172,7 +172,7 @@ function (ipa::Union{IPCrossA, IPA})(TiL::Tuple{AbstractArray,AbstractArray}, si
 
     # Setting mask to the correct dim for broadcasting. 
     if mask != 0 
-        mask = reshape(mask, 1, N_frames_R, N_frames_L, 1) 
+        mask = unsqueeze(mask, dims = 1) 
     end
 
     att = Flux.softmax(w_L .* (att_arg .+ bij) .+ mask, dims = 3)