diff --git a/console_play.py b/console_play.py index d726d0b..a3872a3 100644 --- a/console_play.py +++ b/console_play.py @@ -45,7 +45,8 @@ def play_unconstrained(): if __name__ == '__main__': - play_unconstrained() + import cProfile + cProfile.run('play_unconstrained()') diff --git a/generator/ctrl/ctrl_generator.py b/generator/ctrl/ctrl_generator.py index f823b15..1517be1 100644 --- a/generator/ctrl/ctrl_generator.py +++ b/generator/ctrl/ctrl_generator.py @@ -23,7 +23,7 @@ def loss(labels, logits): class CTRLGenerator(): - def __init__(self, control_code="Apocalypse ", generate_num=28, temperature=0.5, topk=40, nucleus_prob=0, lower_memory=True): + def __init__(self, control_code="Apocalypse ", generate_num=28, temperature=0.5, topk=40, nucleus_prob=0, lower_memory=False): self.generate_num=generate_num model_dir = "generator/ctrl/model/aidungeon2model/"