mypy5

Author: SeBorgey

autointent/modules/scoring/_gcn/gcn_model.py

@@ -1,5 +1,6 @@
 import json
 from pathlib import Path
+from typing import cast
 
 import torch
 from pydantic import BaseModel
@@ -88,7 +89,7 @@ def create_correlation_matrix(
         reweighted_adj = adj_matrix_no_self_loop * weights_p.unsqueeze(1)
         reweighted_adj.fill_diagonal_(1 - p)
 
-        return reweighted_adj
+        return cast(torch.Tensor, reweighted_adj)
 
     def set_correlation_matrix(self, train_labels: torch.Tensor) -> None:
         corr_matrix = self.create_correlation_matrix(
@@ -133,4 +134,4 @@ def load(cls, path: Path, device: str | None = None) -> Self:
 
         instance = instance.to(device)
         instance.eval()
-        return instance
+        return instance