From 171170f506abac954191206b44d6d390f3fb7986 Mon Sep 17 00:00:00 2001
From: Ziyi Yin <ziyi.yin@gatech.edu>
Date: Sun, 8 Jan 2023 11:00:14 -0500
Subject: [PATCH] fix reshape

---
 src/TimeModeling/Types/abstract.jl | 10 ++--------
 1 file changed, 2 insertions(+), 8 deletions(-)

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 `*`