diff --git a/giant/cli.py b/giant/cli.py
index 65c6cac..bdce197 100644
--- a/giant/cli.py
+++ b/giant/cli.py
@@ -300,9 +300,9 @@ def predict(
                 "pre_dir_x":    chunk["pre_dir"][:, 0],
                 "pre_dir_y":    chunk["pre_dir"][:, 1],
                 "pre_dir_z":    chunk["pre_dir"][:, 2],
-                "material_id":  chunk["material"],
+                "material":     chunk["material"],
                 "layer_id":     chunk["layer_id"],
-                "n_secondaries": chunk["n_sec"],
+                "n_sec":         chunk["n_sec"],
                 "step_length":  step_length,
                 "delta_e":      delta_e,
                 "edep":         edep,
diff --git a/giant/data/loader.py b/giant/data/loader.py
index d9aab78..aa7955e 100644
--- a/giant/data/loader.py
+++ b/giant/data/loader.py
@@ -23,9 +23,9 @@ def _df_to_dict(df: pd.DataFrame) -> dict[str, np.ndarray]:
         "pre_pos": df[["pre_x", "pre_y", "pre_z"]].to_numpy(dtype=np.float32),
         "pre_energy": df["pre_energy"].to_numpy(dtype=np.float32),
         "pre_dir": df[["pre_dir_x", "pre_dir_y", "pre_dir_z"]].to_numpy(dtype=np.float32),
-        "material": df["material_id"].to_numpy(dtype=np.int32),
+        "material": df["material"].to_numpy(dtype=np.int32),
         "layer_id": df["layer_id"].to_numpy(dtype=np.int32),
-        "n_sec": df["n_secondaries"].to_numpy(dtype=np.int32),
+        "n_sec": df["child_track_ids"].apply(len).to_numpy(dtype=np.int32),
         "step_length": df["step_length"].to_numpy(dtype=np.float32),
         "delta_e": (df["pre_energy"] - df["post_energy"]).to_numpy(dtype=np.float32),
         "edep": df["edep"].to_numpy(dtype=np.float32),
@@ -53,7 +53,7 @@ _COND_COLS = [
     "event_id", "pdg",
     "pre_x", "pre_y", "pre_z", "pre_energy",
     "pre_dir_x", "pre_dir_y", "pre_dir_z",
-    "material_id", "layer_id", "n_secondaries",
+    "material", "layer_id", "child_track_ids",
 ]
 
 
@@ -64,9 +64,9 @@ def _cond_df_to_dict(df: pd.DataFrame) -> dict[str, np.ndarray]:
         "pre_pos": df[["pre_x", "pre_y", "pre_z"]].to_numpy(dtype=np.float32),
         "pre_energy": df["pre_energy"].to_numpy(dtype=np.float32),
         "pre_dir": df[["pre_dir_x", "pre_dir_y", "pre_dir_z"]].to_numpy(dtype=np.float32),
-        "material": df["material_id"].to_numpy(dtype=np.int32),
+        "material": df["material"].to_numpy(dtype=np.int32),
         "layer_id": df["layer_id"].to_numpy(dtype=np.int32),
-        "n_sec": df["n_secondaries"].to_numpy(dtype=np.int32),
+        "n_sec": df["child_track_ids"].apply(len).to_numpy(dtype=np.int32),
     }
 
 
@@ -91,13 +91,13 @@ def build_index_maps(
 def build_index_maps_from_files(
     files: list[Path],
 ) -> tuple[dict[int, int], dict[int, int]]:
-    """Scan only pdg and material_id columns across all files (2-column read)."""
+    """Scan only pdg and material columns across all files (2-column read)."""
     pdg_vals: set[int] = set()
     mat_vals: set[int] = set()
     for path in files:
-        df = pd.read_parquet(path, columns=["pdg", "material_id"])
+        df = pd.read_parquet(path, columns=["pdg", "material"])
         pdg_vals.update(int(v) for v in df["pdg"].unique())
-        mat_vals.update(int(v) for v in df["material_id"].unique())
+        mat_vals.update(int(v) for v in df["material"].unique())
     return (
         {v: i for i, v in enumerate(sorted(pdg_vals))},
         {v: i for i, v in enumerate(sorted(mat_vals))},