From 746e3840b86a7ec18e12415d4496bb02a2d3cc88 Mon Sep 17 00:00:00 2001
From: Christian Locatelli <christian.locatelli@gmx.de>
Date: Mon, 29 Jun 2026 16:49:47 +0200
Subject: [PATCH] Change minor things to run example files on nubench dataset

---
 examples/04_training/01_train_dynedge.py  | 2 +-
 src/graphnet/data/curated_datamodule.py   | 8 ++++----
 src/graphnet/datasets/nubench_datasets.py | 2 +-
 3 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/examples/04_training/01_train_dynedge.py b/examples/04_training/01_train_dynedge.py
index 76df46756..cf45f11a3 100644
--- a/examples/04_training/01_train_dynedge.py
+++ b/examples/04_training/01_train_dynedge.py
@@ -12,7 +12,7 @@
 from graphnet.models import StandardModel
 from graphnet.models.detector.prometheus import Prometheus
 from graphnet.models.gnn import DynEdge
-from graphnet.models.graphs import KNNGraph
+from graphnet.models.data_representation import KNNGraph
 from graphnet.models.task.reconstruction import EnergyReconstruction
 from graphnet.training.callbacks import PiecewiseLinearLR
 from graphnet.training.loss_functions import LogCoshLoss
diff --git a/src/graphnet/data/curated_datamodule.py b/src/graphnet/data/curated_datamodule.py
index 9194e38c5..e96d9505e 100644
--- a/src/graphnet/data/curated_datamodule.py
+++ b/src/graphnet/data/curated_datamodule.py
@@ -25,7 +25,7 @@ class CuratedDataset(GraphNeTDataModule):
 
     Curated Datasets in GraphNeT are pre-converted datasets that have
     been prepared for training and evaluation of deep learning models.
-    On these Datasets, graphnet users can train and benchmark their
+    On these Datasets, GraphNeT users can train and benchmark their
     models against SOTA methods.
     """
 
@@ -325,9 +325,9 @@ def prepare_data(self) -> None:
             # Download, unzip and delete zipped file
             os.makedirs(self.dataset_dir)
             file_path = os.path.join(self.dataset_dir, file_hash + ".tar.gz")
-            os.system(f"wget -O {file_path} {self._mirror}/{file_hash}")
-            os.system(f"tar -xf {file_path} -C {self.dataset_dir}")
-            os.system(f"rm {file_path}")
+            os.system(f"wget -O '{file_path}' {self._mirror}/{file_hash}")
+            os.system(f"tar -xf '{file_path}' -C {self.dataset_dir}")
+            os.system(f"rm '{file_path}'")
 
 
 class IceCubeHostedDataset(CuratedDataset):
diff --git a/src/graphnet/datasets/nubench_datasets.py b/src/graphnet/datasets/nubench_datasets.py
index 3465ac5ce..c93f436ed 100644
--- a/src/graphnet/datasets/nubench_datasets.py
+++ b/src/graphnet/datasets/nubench_datasets.py
@@ -24,8 +24,8 @@
     "sensor_pos_x",
     "sensor_pos_y",
     "sensor_pos_z",
-    "charge",
     "t",
+    "charge",
 ]
 
 TRUTH_NUBENCH = [