Fix creation of int tensors

Author: ValerianRey

src/torchjd/sparse/_structured_sparse_tensor.py

@@ -28,7 +28,7 @@ def __new__(cls, physical: Tensor, strides: Tensor):
         # (which is bad!)
         assert not physical.requires_grad or not torch.is_grad_enabled()
 
-        pshape = torch.tensor(physical.shape)
+        pshape = tensor(physical.shape, dtype=torch.int64)
         vshape = strides @ (pshape - 1) + 1
         return Tensor._make_wrapper_subclass(
             cls, tuple(vshape.tolist()), dtype=physical.dtype, device=physical.device
@@ -186,7 +186,7 @@ def strides_v2(p_dims: list[int], physical_shape: list[int]) -> list[int]:
 
 
 def get_groupings(pshape: list[int], strides: Tensor) -> list[list[int]]:
-    strides_time_pshape = strides * tensor(pshape)
+    strides_time_pshape = strides * tensor(pshape, dtype=torch.int64)
     groups = {i: {i} for i, column in enumerate(strides.T)}
     group_ids = [i for i in range(len(strides.T))]
     for i1, i2 in itertools.combinations(range(strides.shape[1]), 2):
@@ -260,15 +260,15 @@ def get_full_source(source: list[int], destination: list[int], ndim: int) -> lis
     """
 
     idx = torch.full((ndim,), -1, dtype=torch.int64)
-    idx[destination] = tensor(source)
+    idx[destination] = tensor(source, dtype=torch.int64)
     source_set = set(source)
-    idx[idx.eq(-1)] = tensor([i for i in range(ndim) if i not in source_set])
+    idx[idx.eq(-1)] = tensor([i for i in range(ndim) if i not in source_set], dtype=torch.int64)
 
     return idx.tolist()
 
 
 def fix_dim_of_size_1(physical: Tensor, strides: Tensor) -> tuple[Tensor, Tensor]:
-    is_of_size_1 = torch.tensor([s == 1 for s in physical.shape], dtype=torch.bool)
+    is_of_size_1 = tensor([s == 1 for s in physical.shape], dtype=torch.bool)
     return physical.squeeze(), strides[:, ~is_of_size_1]