fix get_n_best_policies()

autoaug/autoaugment_learners/AaLearner.py

@@ -476,8 +476,8 @@ class AaLearner:
             list of best n policies
         """
 
-        number_policies = max(number_policies, len(self.history))
+        number_policies = min(number_policies, len(self.history))
 
         inter_pol = sorted(self.history, key=lambda x: x[1], reverse = True)[:number_policies]
 
-        return inter_pol[:number_policies]
+        return inter_pol

autoaug/autoaugment_learners/UcbLearner.py

@@ -251,11 +251,11 @@ class UcbLearner(RsLearner):
 
         temp_history = list(zip(self.policies, temp_avg_accs))
 
-        number_policies = max(number_policies, len(temp_history))
+        number_policies = min(number_policies, len(temp_history))
 
         inter_pol = sorted(temp_history, key=lambda x: x[1], reverse = True)[:number_policies]
 
-        return inter_pol[:number_policies]
+        return inter_pol