example with a model

ecmwf · Oct 8, 2024 · 63cd66d · 63cd66d
1 parent d3bf018
commit 63cd66d
Showing 1 changed file with 5 additions and 5 deletions.
diff --git a/src/anemoi/models/preprocessing/normalizer.py b/src/anemoi/models/preprocessing/normalizer.py
@@ -11,7 +11,7 @@
 import warnings
 from typing import Optional
 
-from anemoi.datasets import AnemoiState
+from anemoi.utils.data_structures import NumpyNestedAnemoiTensor
 
 import numpy as np
 import torch
@@ -53,10 +53,10 @@ def __init__(
         mean = statistics["mean"]
         stdev = statistics["stdev"]
 
-        mean = AnemoiState(mean)
-        stdev = AnemoiState(stdev)
-        minimum = AnemoiState(minimum)
-        maximum = AnemoiState(maximum)
+        mean = NumpyNestedAnemoiTensor(mean)
+        stdev = NumpyNestedAnemoiTensor(stdev)
+        minimum = NumpyNestedAnemoiTensor(minimum)
+        maximum = NumpyNestedAnemoiTensor(maximum)
 
         self._validate_normalization_inputs(name_to_index_training_input, minimum, maximum, mean, stdev)