diff --git a/whereami/get_data.py b/whereami/get_data.py
index 800bfd4..467d617 100644
--- a/whereami/get_data.py
+++ b/whereami/get_data.py
@@ -38,5 +38,5 @@ def get_train_data(folder=None):
                 for line in f:
                     data.append(json.loads(line))
             X.extend(data)
-            y.extend([fname.removesuffix(".txt")] * len(data))
+            y.extend([os.path.splitext(fname)[0]] * len(data))
     return X, y