FreezeEncoder

class aixd.mlmodel.architecture.two_stage_model.FreezeEncoder(*args: Any, **kwargs: Any)[source]

Bases: BaseFinetuning

A PyTorch Lightning callback class that extends the BaseFinetuning class and freezes the encoder of the CondAEModel before training.

Methods

finetune_function

freeze_before_training

Freezes the encoder of the CondAEModel before training.