| import torch | |
| import torch.nn as nn | |
| from transformers import PreTrainedModel | |
| from .aes import VAE | |
| from .cceVAEConfig import cceVAEConfig | |
| class cceVAE(PreTrainedModel): | |
| config_class = cceVAEConfig | |
| def __init__(self, config): | |
| super().__init__(config) | |
| self.model = VAE( | |
| d=config.d, | |
| input_size=config.input_size, | |
| z_dim=config.z_dim, | |
| fmap_sizes=config.fmap_sizes, | |
| to_1x1=config.to_1x1, | |
| conv_params=config.conv_params, | |
| tconv_params=config.tconv_params, | |
| normalization_op=config.normalization_op, | |
| normalization_params=config.normalization_params, | |
| activation_op=config.activation_op, | |
| activation_params=config.activation_params, | |
| block_op=config.block_op, | |
| block_params=config.block_params) | |
| def forward(self, x): | |
| return self.model(x) |