Frameowrk dependant float conversion

huggingface · Jun 4, 2024 · 65c0570 · 65c0570
1 parent 489c63d
commit 65c0570
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/src/transformers/pipelines/text_classification.py b/src/transformers/pipelines/text_classification.py
@@ -202,7 +202,12 @@ def postprocess(self, model_outputs, function_to_apply=None, top_k=1, _legacy=Tr
                 function_to_apply = ClassificationFunction.NONE
 
         outputs = model_outputs["logits"][0]
-        outputs = outputs.float().numpy()
+
+        if self.framework == "pt":
+            # To enable using fp16 and bf16
+            outputs = outputs.float().numpy()
+        else:
+            outputs = outputs.numpy()
 
         if function_to_apply == ClassificationFunction.SIGMOID:
             scores = sigmoid(outputs)