mypy3

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
@@ -100,9 +101,9 @@ def forward(self, bert_features: torch.Tensor, label_embeddings: torch.Tensor) -
         classifiers = label_embeddings
         for i in range(len(self.gcn_layers)):
             classifiers = self.gcn_layers[i](self.correlation_matrix, classifiers)
-            classifiers = self.activations[i](classifiers)
+            classifiers = self.activations[i](classifiers)  # type: ignore[operator]
 
-        return torch.matmul(bert_features, classifiers.T)
+        return torch.matmul(bert_features, cast(torch.Tensor, classifiers).T)
 
     def dump(self, path: Path) -> None:
         metadata = GCNModelDumpMetadata(