From de92448ac1df0106ad56dc1c7338a355deff4e18 Mon Sep 17 00:00:00 2001
From: Jirka <jirka.borovec@seznam.cz>
Date: Wed, 9 Aug 2023 10:43:57 +0200
Subject: [PATCH] fixing `_strip` for eventually int

---
 src/flash/core/data/utilities/classification.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/src/flash/core/data/utilities/classification.py b/src/flash/core/data/utilities/classification.py
index 19a40e0449..6bd6992a5a 100644
--- a/src/flash/core/data/utilities/classification.py
+++ b/src/flash/core/data/utilities/classification.py
@@ -42,8 +42,11 @@ def _as_list(x: Union[List, Tensor, np.ndarray]) -> List:
     return x
 
 
-def _strip(x: str) -> str:
-    return x.strip(", ")
+def _strip(x: Union[str, int]) -> str:
+    """Replace both ` ` and `,` from str."""
+    if isinstance(x, str):
+        return x.strip(", ")
+    return str(x)
 
 
 @dataclass