- Fix CUDA device mismatch in WaveletConv and ledoit_wolf

- WaveletConv.forward hardcoded self.device (defaults to cpu), ignoring the input tensor's device. ledoit_wolf didn't coerce shrink_mat to the
input device/dtype. Both caused RuntimeError when running on gpu.

Author: timnaher

spd_learn/functional/regularize.py

@@ -79,7 +79,7 @@ def ledoit_wolf(
     alpha = torch.sigmoid(shrinkage)
 
     return (1 - alpha)[..., None, None] * covariances + alpha[..., None, None] * (
-        mu[..., None] * shrink_mat
+        mu[..., None] * shrink_mat.to(device=covariances.device, dtype=covariances.dtype)
     )
 
 

spd_learn/modules/wavelet.py

@@ -203,4 +203,4 @@ def forward(self, X: Tensor) -> Tensor:
             n_batch, n_freqs, n_sensors, n_epochs, n_times = X_conv.shape
             X_conv = X_conv.view(n_batch, n_freqs, n_sensors, n_epochs * n_times)
 
-        return X_conv.to(device=self.device, dtype=self.dtype)
+        return X_conv.to(device=X.device, dtype=self.dtype)