def GlowTrainAndDecode(self, in_frames=1, out_frames=1,
                         latent_dist_encoder="pointwise",
                         gen_mode="conditional", pretrain_steps=-1,
                         num_train_frames=-1, cond_first_frame=False,
                         apply_dilations=False, activation="relu"):
    """Test 1 forward pass and sampling gives reasonable results."""
    if num_train_frames == -1:
      total_frames = in_frames + out_frames
    else:
      total_frames = num_train_frames

    curr_dir = tempfile.mkdtemp()
    model_path = os.path.join(curr_dir, "model")

    # Training pipeline
    with tf.Graph().as_default():
      hparams = next_frame_glow.next_frame_glow_hparams()
      hparams = fill_hparams(hparams, in_frames, out_frames,
                             gen_mode, latent_dist_encoder, pretrain_steps,
                             num_train_frames, cond_first_frame,
                             apply_dilations, activation)
      features = create_basic_features(hparams)
      model = next_frame_glow.NextFrameGlow(hparams, MODES.TRAIN)
      _, train_op = model(features)
      if self.should_run_session(hparams):
        self.RunModel(model, train_op, hparams, features, total_frames,
                      model_path)

    # Inference pipeline
    with tf.Graph().as_default():
      hparams = next_frame_glow.next_frame_glow_hparams()
      if hparams.gen_mode == "unconditional":
        hparams.video_num_target_frames = 1
      hparams = fill_hparams(hparams, in_frames, out_frames,
                             gen_mode, latent_dist_encoder, pretrain_steps,
                             num_train_frames, cond_first_frame,
                             apply_dilations, activation)
      features = create_basic_features(hparams)
      model = next_frame_glow.NextFrameGlow(
          hparams, tf_estimator.ModeKeys.PREDICT)
      predictions = model.infer(features)
      outputs = predictions["outputs"]
      model_path = os.path.join(curr_dir, "model")

      if self.should_run_session(hparams):
        with tf.Session() as session:
          saver = tf.train.Saver()
          saver.restore(session, model_path)
          outputs_np = session.run(outputs)
          self.assertEqual(outputs_np.shape, (16, out_frames, 64, 64, 3))
          self.assertTrue(np.all(outputs_np <= 255))
          self.assertTrue(np.all(outputs_np >= 0))
Test 1 forward pass and sampling gives reasonable results.
tensor2tensor/models/video/nfg_test_utils.py