Обучение экстрактора пользовательских функций в стабильных базовых условиях3. Начиная с предварительно обученных весов?

Я использую следующий экстрактор пользовательских функций для своей модели StableBaselines3:

import torch.nn as nn
from stable_baselines3 import PPO

class Encoder(nn.Module):
    def __init__(self, input_dim, embedding_dim, hidden_dim, output_dim=2):
        super(Encoder, self).__init__()
        self.encoder = nn.Sequential(
            nn.Linear(input_dim, embedding_dim),
            nn.ReLU()
        )
        self.regressor = nn.Sequential(
            nn.Linear(embedding_dim, hidden_dim),
            nn.ReLU(),
        )
    
    def forward(self, x):
        x = self.encoder(x)
        x = self.regressor(x)
        return x
    
model = Encoder(input_dim, embedding_dim, hidden_dim)
model.load_state_dict(torch.load('trained_model.pth'))

# Freeze all layers
for param in model.parameters():
    param.requires_grad = False

class CustomFeatureExtractor(BaseFeaturesExtractor):
    def __init__(self, observation_space, features_dim):
        super(CustomFeatureExtractor, self).__init__(observation_space, features_dim)
        self.model = model  # Use the pre-trained model as the feature extractor

        self._features_dim = features_dim

    def forward(self, observations):
        features = self.model(observations)
        return features

policy_kwargs = {
        "features_extractor_class": CustomFeatureExtractor,
        "features_extractor_kwargs": {"features_dim": 64}
    }

 model = PPO("MlpPolicy", env=envs, policy_kwargs=policy_kwargs)

На данный момент модель хорошо обучена, без проблем и с хорошими результатами. Теперь я хочу не замораживать веса и попытаться обучить Feature Extractor, начиная с начального предварительно обученного веса. Как я могу сделать это с помощью такого специального экстрактора функций, определенного как класс внутри другого класса? Мой экстрактор функций отличается от описанного в документации , поэтому я не уверен, что он будет обучен. Или он начнет обучение, если я разморозлю слои?