diff --git a/huracanpy/_data/_netcdf.py b/huracanpy/_data/_netcdf.py
index dc71009..0931aaf 100644
--- a/huracanpy/_data/_netcdf.py
+++ b/huracanpy/_data/_netcdf.py
@@ -24,11 +24,13 @@ def load(filename, **kwargs):
     for npoints, tr_id in zip(rowsize.data, trajectory_id.data):
         trajectory_id_stretched.extend([tr_id] * npoints)
 
+    dataset = dataset.drop_vars([trajectory_id.name, rowsize.name])
+
     dataset["track_id"] = (sample_dimension, trajectory_id_stretched)
     # Keep attributes (including cf_role)
     dataset["track_id"].attrs = trajectory_id.attrs
 
-    return dataset.drop_vars([trajectory_id.name, rowsize.name])
+    return dataset
 
 
 def save(dataset, filename):
diff --git a/tests/test_huracanpy.py b/tests/test_huracanpy.py
index 82b4b0c..51f51c3 100644
--- a/tests/test_huracanpy.py
+++ b/tests/test_huracanpy.py
@@ -64,7 +64,7 @@ def test_save(filename, tracker, extension, tmp_path):
     # Reload the data and check it is still the same
     data_ = huracanpy.load(str(tmp_path / f"tmp_file.{extension}"))
 
-    for var in list(data_.variables) + list(data_.coords):
+    for var in list(data.variables) + list(data.coords):
         # Work around for xarray inconsistent loading the data as float or double
         # depending on fill_value and scale_factor
         # np.testing.assert_allclose doesn't work for datetime64