add type hints

NVlabs · Sep 8, 2024 · 397a460 · 397a460
1 parent c4c47a9
commit 397a460
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/earth2grid/_regrid.py b/earth2grid/_regrid.py
@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import math
+from typing import Dict
 
 import einops
 import netCDF4 as nc
@@ -43,7 +44,7 @@ def forward(self, z):
         return output
 
     @staticmethod
-    def from_state_dict(d):
+    def from_state_dict(d: Dict[str, torch.Tensor]) -> "Regridder":
         n, p = d["index"].shape
         regridder = Regridder(n, p)
         regridder.load_state_dict(d)