Fix type error when calling create_dataarray_from_tensor (#106)

mllam · Feb 6, 2025 · 82939a2 · 82939a2
1 parent d494cd3
commit 82939a2
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -18,6 +18,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Fixed
 - Only print on rank 0 to avoid duplicates of all print statements.
 [\#103](https://github.com/mllam/neural-lam/pull/103) @simonkamuk @sadamov
+- Fix duplicate tensor copy to CPU [\#106](https://github.com/mllam/neural-lam/pull/106) @observingClouds
 
 ### Maintenance
 - update ci/cd testing setup to install torch version compatible with neural-lam

diff --git a/neural_lam/models/ar_model.py b/neural_lam/models/ar_model.py
@@ -181,7 +181,7 @@ def _create_dataarray_from_tensor(
         weather_dataset = WeatherDataset(datastore=self._datastore, split=split)
         time = np.array(time.cpu(), dtype="datetime64[ns]")
         da = weather_dataset.create_dataarray_from_tensor(
-            tensor=tensor.cpu().numpy(), time=time, category=category
+            tensor=tensor, time=time, category=category
         )
         return da