diff --git a/src/TimeModeling/Types/abstract.jl b/src/TimeModeling/Types/abstract.jl
index 7227c00a4..be848d848 100644
--- 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 `*`