diff --git a/tour_model_eval/build_save_model.py b/tour_model_eval/build_save_model.py
index 0a42e97..425d71f 100644
--- a/tour_model_eval/build_save_model.py
+++ b/tour_model_eval/build_save_model.py
@@ -110,7 +110,7 @@ def main():
                     # compute unique label sets and their probabilities in one cluster
                     # 'p' refers to probability
                     unique_labels = user_label_df.groupby(user_label_df.columns.tolist()).size().reset_index(name='count')
-                    unique_labels['p'] = [unique_labels.iloc[i]['count'] / sum_trips for i in range(len(unique_labels))]
+                    unique_labels['p'] = unique_labels.count / sum_trips
                     labels_columns = user_label_df.columns.to_list()
                     for i in range(len(unique_labels)):
                         one_set_labels = {}
@@ -165,4 +165,3 @@ def main():
 
 
 
-