Repository: speechbrain/speechbrain
Issue category: Bug Report

Issue:
Weighted SSL model not unfreezable
### Describe the bug

In our HF Weighted SSL model implementation, we `detach()` the hidden states, meaning weights are not updated.

Relevant code:
https://github.com/speechbrain/speechbrain/blob/develop/speechbrain/lobes/models/huggingface_transformers/weighted_ssl.py#L81

```
hidden_states = torch.stack(feats.hidden_states, dim=0).detach()
```

### Expected behaviour

If passing `freeze=False` I'd expect the weights to get updated.

### To Reproduce

_No response_

### Environment Details

_No response_

### Relevant Log Output

_No response_

### Additional Context

_No response_


Hints:
Also if `layernorm=True` then the hidden states are converted to a list which causes a program crash. They should be re-stacked into a tensor.

Patch:
diff --git a/speechbrain/lobes/models/huggingface_transformers/weighted_ssl.py b/speechbrain/lobes/models/huggingface_transformers/weighted_ssl.py
index 1507f85093..2dad9e1e46 100644
--- a/speechbrain/lobes/models/huggingface_transformers/weighted_ssl.py
+++ b/speechbrain/lobes/models/huggingface_transformers/weighted_ssl.py
@@ -42,6 +42,8 @@ class WeightedSSLModel(HFTransformersInterface):
     freeze : bool (default: True)
         If True, the model is frozen. If False, the model will be trained
         alongside with the rest of the pipeline.
+    **kwargs : dict
+        Additional arguments to pass to HFTransformersInterface
 
     Example
     -------
@@ -52,14 +54,19 @@ class WeightedSSLModel(HFTransformersInterface):
     >>> outputs = model(inputs)
     """
 
-    def __init__(self, hub, save_path="", layernorm=False, freeze=False):
-        super().__init__(source=hub, save_path=save_path, freeze=freeze)
+    def __init__(
+        self, hub, save_path="", layernorm=False, freeze=False, **kwargs
+    ):
+        super().__init__(
+            source=hub, save_path=save_path, freeze=freeze, **kwargs
+        )
         self.model.eval()
+        self.layernorm = layernorm
+        self.freeze = freeze
         self.num_layers = self.config.num_hidden_layers + 1
         # Initializing the learnable weights
         zero_init = torch.cat([torch.zeros(self.num_layers)])
         self.weights = torch.nn.Parameter(zero_init, requires_grad=True)
-        self.layernorm = layernorm
 
     def forward(self, wav, wav_lens=None):
         """This method outputs a weighted sum of the layer representations of the SSL encoder
@@ -78,21 +85,25 @@ def forward(self, wav, wav_lens=None):
         """
 
         feats = self.model(wav)
-        hidden_states = torch.stack(feats.hidden_states, dim=0).detach()
+        if self.freeze:
+            hidden_states = torch.stack(feats.hidden_states, dim=0).detach()
+        else:
+            hidden_states = torch.stack(feats.hidden_states, dim=0)
+
         # First dimension should be equal to the number of layers in the hparams
         assert (
             self.num_layers == hidden_states.shape[0]
         ), "Num layers not equal to num hidden states"
-        norm_weights = torch.nn.functional.softmax(self.weights, dim=-1)
+
         # Layernorming the layers representations if asked
         if self.layernorm:
-            hidden_states = [
-                F.layer_norm(t, (t.shape[-1],)) for t in hidden_states
-            ]
+            normalized_shape = (hidden_states.size(-1),)
+            hidden_states = F.layer_norm(hidden_states, normalized_shape)
+
         # Summing the weighted layers
-        weighted_feats = (
-            hidden_states * norm_weights[:, None, None, None]
-        ).sum(axis=0)
+        norm_weights = F.softmax(self.weights, dim=-1).view(-1, 1, 1, 1)
+        weighted_feats = (hidden_states * norm_weights).sum(axis=0)
+
         return weighted_feats
 
     def override_config(self, config):


Test patch: