Fix shape descriptions in calculate_loss method

[yuta0x89]

The content of latents is defined in the following code in modeling_sd_base.py:

latents = self.prepare_latents(
    batch_size * num_images_per_prompt,
    num_channels_latents,
    height,
    width,
    prompt_embeds.dtype,
    device,
    generator,
    latents,
)

As a result, the shape of latents and next_latents mentioned in the description of the calculate_loss method in the DDPOTrainer class should be [batch_size, num_channels_latents, height, width].

[lvwerra]

Letting @sayakpaul have a look too.