def extract_features(
        self,
        source: torch.Tensor,
        padding_mask: Optional[torch.Tensor] = None,
        mask: bool = False,
        ret_conv: bool = False,
        output_layer: Optional[int] = None,
        ret_layer_results: bool = False,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """Extract features for only speech input"""
        with torch.no_grad():
            res = self.forward(
                source,
                padding_mask=padding_mask,
                mask=mask,
                features_only=True,
                output_layer=output_layer,
            )
            # {"x": x, "padding_mask": padding_mask, "features": features, "layer_results": layer_results}

            x = res["x"] # B x T x D
            padding_mask = res["padding_mask"]
            if self.add_unit_encoder and (output_layer is None or output_layer > self.cfg.encoder_layers):
                src_tokens, x, _ = self.convert_embeddings(
                    x,
                    padding_mask,
                    mix_with_unit=False,
                    use_pred_unit=False,
                )
                return_all_hiddens=output_layer is not None and output_layer > self.cfg.encoder_layers
                encoder_out = self.unit_encoder(
                    src_tokens,
                    token_embeddings=x,
                    return_all_hiddens=return_all_hiddens,
                )
                res["x"] = encoder_out['encoder_out'][0].transpose(0, 1)  # (B, T, D)
                if return_all_hiddens:
                    res["layer_results"] += encoder_out['encoder_states'][1:1+output_layer-len(res["layer_results"])]
            
            feature = res["features"] if ret_conv else res["x"]
            if ret_layer_results:
                feature = (feature, res["layer_results"])

            return feature, padding_mask
Extract features for only speech input
speechlm/SpeechLM.py