add type hints

pytorch · Sep 12, 2024 · 28d3b4b · 28d3b4b
1 parent 8458ae1
commit 28d3b4b
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/ignite/metrics/hsic.py b/ignite/metrics/hsic.py
@@ -133,6 +133,8 @@ def update(self, output: Sequence[Tensor]) -> None:
         xx = X @ X.T
         rx = xx.diag().unsqueeze(0).expand_as(xx)
         dxx = rx.T + rx - xx * 2
+
+        vx: Union[Tensor, float]
         if self.sigma_x < 0:
             vx = torch.quantile(dxx, 0.5)
         else:
@@ -142,6 +144,8 @@ def update(self, output: Sequence[Tensor]) -> None:
         yy = Y @ Y.T
         ry = yy.diag().unsqueeze(0).expand_as(yy)
         dyy = ry.T + ry - yy * 2
+
+        vy: Union[Tensor, float]
         if self.sigma_y < 0:
             vy = torch.quantile(dyy, 0.5)
         else: