diff --git a/SyntheticControlMethods/main.py b/SyntheticControlMethods/main.py
index aa7a306..0d96cf9 100644
--- a/SyntheticControlMethods/main.py
+++ b/SyntheticControlMethods/main.py
@@ -458,13 +458,13 @@ def difference_data(self, dataset, not_diff_cols):
         ignore_all_cols = not_diff_cols == None
 
         #Compute difference of outcome variable
-        modified_dataset[data.outcome_var] = modified_dataset.groupby(data.id)[data.outcome_var].apply(
+        modified_dataset[data.outcome_var] = modified_dataset.groupby(data.id, group_keys=False)[data.outcome_var].apply(
                                             lambda unit: unit.interpolate(method='linear', limit_direction="both")).diff()
         
         #For covariates
         for col in data.covariates:
             #Fill in missing values using unitwise linear interpolation
-            modified_dataset[col] = modified_dataset.groupby(data.id)[col].apply(
+            modified_dataset[col] = modified_dataset.groupby(data.id, group_keys=False)[col].apply(
                                     lambda unit: unit.interpolate(method='linear', limit_direction="both"))
             
             #Compute change from previous period
diff --git a/SyntheticControlMethods/optimize.py b/SyntheticControlMethods/optimize.py
index e264c01..9073db4 100644
--- a/SyntheticControlMethods/optimize.py
+++ b/SyntheticControlMethods/optimize.py
@@ -86,10 +86,10 @@ def optimize(self,
                     v_0 = np.append(v_0, 0)
             else:
                 #Dirichlet distribution returns a valid pmf over n_covariates states
-                v_0 = self.original_data.rng.dirichlet(np.ones(data.n_covariates), size=1)
+                v_0 = self.original_data.rng.dirichlet(np.ones(data.n_covariates))
                 if pen == "auto":
                     #if pen =="auto", we have an additional parameter to optimize over, so we append it
-                    v_0 = np.append(v_0, self.original_data.rng.lognormal(1.5, 1, size=1)) #Still experimenting with what distribution is appropriate
+                    v_0 = np.append(v_0, self.original_data.rng.lognormal(1.5, 1)) #Still experimenting with what distribution is appropriate
             
 
             #Required to have non negative values