mypy

Author: xadupre

onnx_diagnostic/reference/torch_evaluator.py

@@ -135,7 +135,7 @@ def _build_kernels(self, nodes: Sequence[onnx.NodeProto]):
 
     def run(
         self, outputs: Optional[List[str]], feeds: Dict[str, torch.Tensor]
-    ) -> List[torch.Tensor]:
+    ) -> List[Optional[torch.Tensor]]:
         """
         Runs the ONNX model.
 
@@ -182,12 +182,12 @@ def run(
                 self.runtime_info[name].clean_value()
 
         # outputs
-        res = [self.runtime_info[o].value.tensor for o in outputs]
+        res = [self.runtime_info[o].value.tensor for o in outputs]  # type: ignore[assignment, union-attr]
 
         # clean previous execution
         for k in feeds:
             self.runtime_info[k].clean_value()
         for o in outputs:
             self.runtime_info[o].clean_value()
 
-        return res
+        return res  # type: ignore[return-value]