Fix some Mypy errors.

Author: PierreQuinton

src/torchjd/sparse/_aten_function_overrides/shape.py

@@ -207,7 +207,7 @@ def cat_default(tensors: list[Tensor], dim: int) -> Tensor:
         # Such a physical dimension already exists. Note that an alternative implementation would be
         # to simply always add the physical dimension, and squash it if it ends up being not needed.
         physicals = [t.physical for t in tensors_]
-        pdim = indices[0][0]
+        pdim = cast(int, indices[0, 0].item())
         new_strides = ref_tensor.strides
 
     new_physical = aten.cat.default(physicals, dim=pdim)

src/torchjd/sparse/_structured_sparse_tensor.py

@@ -1,5 +1,6 @@
 import itertools
 import operator
+from collections.abc import Callable
 from functools import wraps
 from itertools import accumulate
 from math import prod
@@ -10,7 +11,7 @@
 
 
 class StructuredSparseTensor(Tensor):
-    _HANDLED_FUNCTIONS = dict()
+    _HANDLED_FUNCTIONS = dict[Callable, Callable]()
 
     @staticmethod
     def __new__(cls, physical: Tensor, strides: Tensor):