diff --git a/java/src/main/java/com/logicalclocks/hsfs/engine/TrainingDatasetEngine.java b/java/src/main/java/com/logicalclocks/hsfs/engine/TrainingDatasetEngine.java
index a586aa69be..33d262ef2e 100644
--- a/java/src/main/java/com/logicalclocks/hsfs/engine/TrainingDatasetEngine.java
+++ b/java/src/main/java/com/logicalclocks/hsfs/engine/TrainingDatasetEngine.java
@@ -67,6 +67,7 @@ public void save(TrainingDataset trainingDataset, Dataset<Row> dataset,
     // Update the original object - Hopsworks returns the full location and incremented version
     trainingDataset.setLocation(apiTD.getLocation());
     trainingDataset.setVersion(apiTD.getVersion());
+    trainingDataset.setId(apiTD.getId());
 
     // Build write options map
     Map<String, String> writeOptions =