fixing _strip for eventually int (#1667)

Lightning-Universe · Aug 9, 2023 · b849ad2 · b849ad2
1 parent e76a61a
commit b849ad2
Showing 1 changed file with 5 additions and 2 deletions.
diff --git 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