def squeeze_video(self, video, init=False):
    """Squeeze a 5-D Tensor video with one timestep to a 4-D frame."""
    if init:
      batch_size = self.hparams.init_batch_size
    else:
      batch_size = self.hparams.batch_size
    frame_shape = [batch_size] + self.hparams.problem.frame_shape
    return tf.reshape(video, frame_shape)
Squeeze a 5-D Tensor video with one timestep to a 4-D frame.
tensor2tensor/models/video/next_frame_glow.py