Fix overly long line

unitaryai · Jan 5, 2025 · 878bf90 · 878bf90
1 parent a669e06
commit 878bf90
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/detoxify/detoxify.py b/detoxify/detoxify.py
@@ -119,8 +119,11 @@ def predict(self, text):
         results = {}
         for i, cla in enumerate(self.class_names):
             results[cla] = (
-                # If the input is a single text, squeezing will remove the dimensionality from the tensor - so `.tolist()` will return a number instead. Otherwise, we'll get the list of scores of that class.
-                scores[:,i].squeeze().tolist()
+                # If the input is a single text, squeezing will remove the dimensionality from the tensor -
+                # so `.tolist()` will return a number instead. Otherwise, we'll get the list of scores of that class.
+                scores[:, i]
+                .squeeze()
+                .tolist()
             )
         return results