File size: 549 Bytes
cd4cfe7 513885e cd4cfe7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 |
import huggingface_hub
from .config import Config
from transformers import PreTrainedModel
from .dcgan import Generator
# config = Config()
# config.save_pretrained("WGAN-GP")
class WGAN_GP(PreTrainedModel):
config_class = Config
def __init__(self, config):
super().__init__(config)
self.generator=Generator(config.cfg["imsize"],config.cfg["img_ch"],config.cfg["zdim"],
config.cfg["norm_type"]["g"],config.cfg["final_activation"]["g"])
def forward(self, input):
return self.generator(input)
|