Fix tensor device assignment in windowed weighted sample covariance a…

…nd update dimension indexing in SoftmaxND classes

src/landmarker/heatmap/decoder.py

@@ -752,7 +752,7 @@ def windowed_weigthed_sample_cov(
                 ]
                 .unsqueeze(0)
                 .unsqueeze(0),
-                torch.tensor([[[window, window]]], dtype=torch.float),
+                torch.tensor([[[window, window]]], dtype=torch.float).to(heatmap.device),
                 spatial_dims=spatial_dims,
                 activation=activation,
             )

src/landmarker/models/utils.py

@@ -5,7 +5,7 @@
 class SoftmaxND(nn.Module):
     def __init__(self, spatial_dims):
         super().__init__()
-        self.dim = (2, 3) if spatial_dims == 2 else (2, 3, 4)
+        self.dim = (-2, -1) if spatial_dims == 2 else (-3, -2, -2)
 
     def forward(self, x):
         out = torch.exp(x)
@@ -15,7 +15,7 @@ def forward(self, x):
 class LogSoftmaxND(nn.Module):
     def __init__(self, spatial_dims):
         super().__init__()
-        self.dim = (2, 3) if spatial_dims == 2 else (2, 3, 4)
+        self.dim = (-2, -1) if spatial_dims == 2 else (-3, -2, -1)
         self.spatial_dims = spatial_dims
 
     def forward(self, x):