How to save hyperparameters of objects passed to LightningModule without breaking checkpoint_loading

[codingS3b]

I'm wondering if there is a good way within a LightningModule to save attributes of an object passed during construction. So far, I'm able to achieve it but problems occur when trying to load models back from a checkpoint.

Hopefully the below example will make my issue more clear.

import pytorch_lightning as pl
import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader, TensorDataset

class FeatureExtractor(torch.nn.Module):
    # an object of this class will be plugged into the lightning module later
    def __init__(self, hidden_dim):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.fc1 = torch.nn.Linear(self.hidden_dim, 2)
    
    def forward(self, x):
        x = self.fc1(x)
        return x

class ModelForTraining(pl.LightningModule):
    def __init__(self, feature_extractor):
        super().__init__()
        self.model = feature_extractor
        
       # NOTE: Here I want to log attributes of the the feature_extractor instance that
       # was passed in so I can see them later e.g. in tensorboard

        # this is fine, but if we call self.hparams, the whole module shows up
        # self.save_hyperparameters()

        # this works at first, but causes problems when loading from checkpoints
        self.save_hyperparameters({'hidden_dim': self.model.hidden_dim})
        
    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr = 1e-3)
        return optimizer
        
    def forward(self, x):
        return self.model(x)

    def training_step(self, batch, batch_idx):
        x, y = batch
        return F.mse_loss(self.forward(x), y)

# create data
hidden_dim = 10
dl = DataLoader(TensorDataset(torch.randn(1000, hidden_dim), 
                              torch.randn(1000, 2)), 
                batch_size = 10)
feature_extractor = FeatureExtractor(hidden_dim=hidden_dim)
model = ModelForTraining(feature_extractor=feature_extractor)

print(model.hparams)

# train model
checkpoint_callback = pl.callbacks.ModelCheckpoint()
trainer = pl.Trainer(gpus=0, max_epochs=1,
                              callbacks=[checkpoint_callback])
trainer.fit(model, dl)

# try to load from checkpoint
# this fails with "__init__() missing 1 required positional argument: 'feature_extractor'"
model_back = ModelForTraining.load_from_checkpoint(checkpoint_callback.best_model_path)