Fix data directory in test

metatensor · Nov 30, 2023 · 4cca96c · 4cca96c
1 parent 366f39c
commit 4cca96c
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/tests/data.py b/tests/data.py
@@ -1,3 +1,4 @@
+import os
 import torch
 
 from metatensor_models.utils.data import Dataset, collate_fn, read_structures, read_targets
@@ -6,8 +7,10 @@
 def test_dataset():
     """Tests the readers and the dataset class."""
 
-    structures = read_structures("data/qm9_reduced_100.xyz")
-    targets = read_targets("data/qm9_reduced_100.xyz", "U0")
+    dataset_path = os.path.join(os.path.dirname(__file__), "data/qm9_reduced_100.xyz")
+
+    structures = read_structures(dataset_path)
+    targets = read_targets(dataset_path, "U0")
 
     dataset = Dataset(structures, targets)
     dataloader = torch.utils.data.DataLoader(dataset, batch_size=10, collate_fn=collate_fn)