black

opentensor · Oct 2, 2023 · e2e348f · e2e348f
1 parent 2d2daba
commit e2e348f
Showing 1 changed file with 6 additions and 6 deletions.
diff --git a/prompting/validators/reward/diversity.py b/prompting/validators/reward/diversity.py
@@ -130,9 +130,9 @@ def regularise(rewards):
 
         # Reward to be at the bottom_k smallest of the 1 - similarity score.
         bottom_k = min(self.history_reward_bottom_k, len(similarity))
-        rewards = torch.topk(
-            (1 - torch.abs(similarity)), bottom_k, largest=False
-        )[0][:, -1]
+        rewards = torch.topk((1 - torch.abs(similarity)), bottom_k, largest=False)[0][
+            :, -1
+        ]
 
         return regularise(rewards)
 
@@ -146,9 +146,9 @@ def regularise(rewards):
 
         # Reward to be at the 10% quantile of the 1 - similarity score.
         bottom_k = min(self.reward_bottom_k, len(similarity))
-        rewards = torch.topk(
-            (1 - torch.abs(similarity)), bottom_k, largest=False
-        )[0][:, -1]
+        rewards = torch.topk((1 - torch.abs(similarity)), bottom_k, largest=False)[0][
+            :, -1
+        ]
 
         return regularise(rewards)