Fix loader for the M5 dataset (#3151)

awslabs · Mar 24, 2024 · 722b9e6 · 722b9e6
1 parent e796bc4
commit 722b9e6
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/src/gluonts/dataset/repository/_m5.py b/src/gluonts/dataset/repository/_m5.py
@@ -61,6 +61,7 @@ def generate_m5_dataset(
             "d",
         ],
         axis=1,
+        errors="ignore",
     )
     cal_features["event_type_1"] = cal_features["event_type_1"].apply(
         lambda x: 0 if str(x) == "nan" else 1
@@ -112,9 +113,11 @@ def generate_m5_dataset(
     ]
 
     # Build target series
-    train_ids = sales_train_validation["id"]
+    train_ids = sales_train_validation["item_id"]
     train_df = sales_train_validation.drop(
-        ["id", "item_id", "dept_id", "cat_id", "store_id", "state_id"], axis=1
+        ["id", "item_id", "dept_id", "cat_id", "store_id", "state_id"],
+        axis=1,
+        errors="ignore",
     )
     test_target_values = train_df.values.copy()
     train_target_values = [ts[:-prediction_length] for ts in train_df.values]