//  [python]
def body(self, features):
    hparams = self.hparams
    input_shape = common_layers.shape_list(features['inputs'])
    batch_size, _, frame_width, frame_height, frame_channels = input_shape  # pylint: disable=unused-variable

    # Swap time and batch axes.
    input_frames = common_video.swap_time_and_batch_axes(
        tf.to_float(features['inputs']))
    target_frames = common_video.swap_time_and_batch_axes(features['targets'])

    # Get actions if exist otherwise use zeros
    input_actions = self.get_input_if_exists(
        features, 'input_action', batch_size, hparams.video_num_input_frames)
    target_actions = self.get_input_if_exists(
        features, 'target_action', batch_size, hparams.video_num_target_frames)

    # Get rewards if exist otherwise use zeros
    # TODO(blazej) enable rewards.
    # input_rewards = self.get_input_if_exists(
    #     features, 'input_reward', batch_size, hparams.video_num_input_frames)
    # target_rewards = self.get_input_if_exists(
    #     features, 'target_reward', batch_size,hparams.video_num_target_frames)
    # all_rewards = tf.concat([input_rewards, target_rewards], axis=0)

    all_actions = tf.concat([input_actions, target_actions], axis=0)
    # flatten actions tensor to have the shape: framesXbatch_sizeXaction_dims.
    actions_shape = common_layers.shape_list(all_actions)
    all_actions = tf.reshape(
        all_actions,
        [actions_shape[0], -1,
         reduce(lambda x, y: x * y, actions_shape[2:])])
    all_frames = tf.concat([input_frames, target_frames], axis=0)

    all_frames = tf.unstack(all_frames, axis=0)
    all_actions = tf.unstack(all_actions, axis=0)

    # TODO(blazej) - most likely this downsize is too strong.
    all_frames = [
        tf.image.resize_images(
            image, (IMG_HEIGHT, IMG_WIDTH),
            method=tf.image.ResizeMethod.BICUBIC)
        for image in all_frames
    ]

    enc_out_all, pred_out_all, _, van_on_enc_all = construct_model(
        all_frames,
        all_actions,
        context_frames=hparams.context_frames,
        hparams=hparams,
        is_training=self.is_training)

    enc_pred_loss, _ = calc_loss_psnr(
        enc_out_all[1:],
        pred_out_all,
        'enc_pred_loss',
        hparams=hparams,
        use_l1_loss=hparams.enc_pred_use_l1_loss)

    van_on_enc_loss, _ = calc_loss_psnr(
        van_on_enc_all,
        all_frames[1:],
        'van_on_enc_loss',
        hparams=hparams)

    enc_pred_loss_scale_delay = max(hparams.enc_pred_loss_scale_delay, 1)
    enc_pred_loss_scale = tf.nn.sigmoid(
        (tf.to_float(tf.train.get_or_create_global_step()
                    ) - enc_pred_loss_scale_delay) /
        (enc_pred_loss_scale_delay * .1)) * hparams.enc_pred_loss_scale
    tf.summary.scalar('enc_pred_loss_scale', enc_pred_loss_scale)
    epva_loss = enc_pred_loss * enc_pred_loss_scale + van_on_enc_loss
    tf.summary.scalar('epva_loss', epva_loss)

    predictions = tf.stack(van_on_enc_all)

    if hparams.clip_pixel_values:
      predictions = tf.clip_by_value(predictions, 0.0, 1.0)

    # TODO(mbz): clean this up!
    def fix_video_dims_and_concat_on_x_axis(x):
      x = tf.transpose(x, [1, 3, 4, 0, 2])
      x = tf.reshape(x, [batch_size, frame_height, frame_channels, -1])
      x = tf.transpose(x, [0, 3, 1, 2])
      return x

    frames_gd = fix_video_dims_and_concat_on_x_axis(target_frames)
    frames_pd = fix_video_dims_and_concat_on_x_axis(predictions)
    side_by_side_video = tf.concat([frames_gd, frames_pd], axis=1)
    tf.summary.image('full_video', side_by_side_video)

    predictions = tf.unstack(predictions)
    predictions = [
        tf.image.resize_images(
            image, (frame_width, frame_height),
            method=tf.image.ResizeMethod.BICUBIC)
        for image in predictions
    ]
    predictions = tf.stack(predictions)

    predictions = common_video.swap_time_and_batch_axes(predictions)
    predictions = tf.slice(predictions,
                           [0, hparams.video_num_input_frames-1, 0, 0, 0],
                           [-1]*5)

    return predictions, {'extra': epva_loss}