KeyError: 'gen_states' when run train.py

[lorafei]

Hi, when I was trying to train savp model like the given example did

CUDA_VISIBLE_DEVICES=0 python scripts/train.py --input_dir data/bair --dataset bair \
  --model savp --model_hparams_dict hparams/bair_action_free/ours_savp/model_hparams.json \
  --output_dir logs/bair_action_free/ours_savp

I got error like this

Traceback (most recent call last):
  File "scripts/train.py", line 354, in <module>
    main()
  File "scripts/train.py", line 174, in main
    model.build_graph(inputs)
  File "/home/jio/vp_ws/video_prediction/models/base_model.py", line 478, in build_graph
    outputs_tuple, losses_tuple, loss_tuple, metrics_tuple = self.tower_fn(self.inputs)
  File "/home/jio/vp_ws/video_prediction/models/base_model.py", line 439, in tower_fn
    g_losses = self.generator_loss_fn(inputs, outputs)
  File "/home/jio/vp_ws/video_prediction/models/base_model.py", line 759, in generator_loss_fn
    gen_states = outputs.get('gen_states_enc', outputs['gen_states'])
KeyError: 'gen_states'

I checked model_hparams.json file theree do exist state_weight param. Any one met this problem?

Thank you!