fix: add missing type hints

IamGianluca · Mar 23, 2024 · 4341f57 · 4341f57
1 parent ceadd8a
commit 4341f57
Showing 1 changed file with 6 additions and 4 deletions.
diff --git a/blazingai/metrics.py b/blazingai/metrics.py
@@ -1,3 +1,5 @@
+from typing import List
+
 import numpy as np
 import torch
 import torchmetrics
@@ -57,10 +59,10 @@ class CrossValMetrics:
     def __init__(self, cfg: DictConfig) -> None:
         self.cfg = cfg
         self.metric = cfg.metric
-        self._trgt = []
-        self._pred = []
-        self._val_scores = []
-        self._trn_scores = []
+        self._trgt: List[torch.Tensor] = []
+        self._pred: List[torch.Tensor] = []
+        self._val_scores: List[torch.Tensor] = []
+        self._trn_scores: List[torch.Tensor] = []
 
     @property
     def trn_metric(self):