From 878bf907901980a83822cc4a502324b9a20d14b0 Mon Sep 17 00:00:00 2001 From: James Thewlis Date: Sun, 5 Jan 2025 14:39:09 +0100 Subject: [PATCH] Fix overly long line --- detoxify/detoxify.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/detoxify/detoxify.py b/detoxify/detoxify.py index c97fa95..489ac28 100644 --- 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