diff --git a/DESCRIPTION b/DESCRIPTION
index abba301..c3f3b9a 100644
--- a/DESCRIPTION
+++ b/DESCRIPTION
@@ -1,8 +1,8 @@
 Package: dtwSat
 Type: Package
 Title: Time-Weighted Dynamic Time Warping for Satellite Image Time Series Analysis
-Version: 1.0.0
-Date: 2023-09-22
+Version: 1.0-1
+Date: 2023-09-25
 Authors@R:
     c(person(given = "Victor",
              family = "Maus",
diff --git a/R/train.R b/R/train.R
index 238e61a..6711ac7 100644
--- a/R/train.R
+++ b/R/train.R
@@ -102,6 +102,10 @@ twdtw_knn1 <- function(x, y, time_weight, cycle_length, time_scale,
     stop(paste("Missing required columns in y:", paste(missing_columns, collapse = ", ")))
   }
 
+  # adjust y column names
+  st_geometry(y) <- 'geom'
+  y <- y[, c(start_column, end_column, label_colum, 'geom')]
+
   # Convert columns to date-time
   y[, start_column] <- to_date_time(y[[start_column]])
   y[, end_column] <- to_date_time(y[[end_column]])