diff --git a/SyntheticControlMethods/optimize.py b/SyntheticControlMethods/optimize.py
index e264c01..d6639bb 100644
--- a/SyntheticControlMethods/optimize.py
+++ b/SyntheticControlMethods/optimize.py
@@ -191,7 +191,7 @@ def total_loss(self, v_0,
             data.in_time_placebo_w = w.value
             
         #Return loss
-        return loss
+        return loss.flatten()
 
     def _get_dsc_outcome(self, w, control_outcome, periods_pre_treatment, treated_pretreatment_outcome):
         '''Method used only by DiffSynth (DSC)
@@ -307,4 +307,4 @@ def random_optimize(self, steps=10**4):
         
         #Store, print, return best solutions
         self.w = best_w
-        return best_w, min_loss
\ No newline at end of file
+        return best_w, min_loss