fix incorrect input type

mllam · Feb 4, 2025 · 06bfb81 · 06bfb81
1 parent f233f87
commit 06bfb81
Showing 1 changed file with 1 addition and 1 deletion.
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