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]