fix reshape

slimgroup · Jan 8, 2023 · 171170f · 171170f
1 parent 6c693a2
commit 171170f
Showing 1 changed file with 2 additions and 8 deletions.
diff --git a/src/TimeModeling/Types/abstract.jl b/src/TimeModeling/Types/abstract.jl
@@ -73,14 +73,8 @@ vec(x::judiMultiSourceVector) = vcat(vec.(x.data)...)
 
 time_sampling(ms::judiMultiSourceVector) = [1 for i=1:ms.nsrc]
 
-function reshape(ms::judiMultiSourceVector, dims::Dims{N}) where N
-    try
-        return reshape(vec(ms), dims)
-    catch e
-        @assert dims[1] == ms.nsrc  ### during AD, size(ms::judiVector) = ms.nsrc
-        return ms
-    end
-end
+reshape(ms::judiMultiSourceVector, dims::Dims{1}) = ms ### during AD, size(ms::judiVector) = ms.nsrc
+reshape(ms::judiMultiSourceVector, dims::Dims{N}) where N = reshape(vec(ms), dims)
 
 ############################################################################################################################
 # Linear algebra `*`