diff --git a/project_1.py b/project_1.py
index ddcd899..5a3d472 100644
--- a/project_1.py
+++ b/project_1.py
@@ -12,7 +12,8 @@
 from torch.distributions import MultivariateNormal
 from torch.utils.data import DataLoader, TensorDataset
 import matplotlib.pyplot as plt
-from sklearn.decomposition import PCA, silhouette_score, adjusted_rand_score
+from sklearn.decomposition import PCA
+from sklearn.metrics import silhouette_score, adjusted_rand_score
 import seaborn as sns
 
 def load_data(data_path, labels_path, dataset_size=10, train=True, standardize=True):
@@ -281,8 +282,8 @@ def main():
         labels_path = r"R:\ENGR_Chon\NIH_Pulsewatch_Database\Adjudication_UConn\final_attemp_4_1_Dong_Ohm"
         saving_path = r"R:\ENGR_Chon\Luis\Research\Casseys_case\Project_1_analysis"
     
-    train_data, segment_names, labels = load_data(data_path, labels_path, dataset_size=141, train=True)
-    test_data, _, _ = load_data(data_path, labels_path, dataset_size=141, train=False)
+    train_data, segment_names, labels = load_data(data_path, labels_path, dataset_size=20, train=True)
+    test_data, _, _ = load_data(data_path, labels_path, dataset_size=30, train=False)
 
     train_dataloader = create_dataloader(train_data)
     test_dataloader = create_dataloader(test_data)
@@ -302,7 +303,7 @@ def main():
     # Perform MFVI for your data
     K = 4  # Number of clusters
     n_optimization_iterations = 100
-    miu_mfvi, pi_mfvi, resp_mfvi = perform_mfvi(train_data, K, n_optimization_iterations, convergence_threshold=1e-5, run_until_convergence=False)
+    miu_mfvi, pi_mfvi, resp_mfvi = perform_mfvi(train_data, K, n_optimization_iterations=1000, convergence_threshold=1e-5, run_until_convergence=False)
     
     # Calculate clustering metrics for MFVI
     zi_mfvi = np.argmax(resp_mfvi, axis=1)