def METHOD_NAME(self, text='', target=None):
    self.unet.train()
    self.unet = self.unet.to(self.device)
    # Convert to latent space
    with torch.no_grad():
        latents = self.vae.encode(
            target.to(dtype=self.weight_dtype)).latent_dist.sample()
    latents = latents * self.vae.config.scaling_factor
    # Sample noise that we'll add to the latents
    noise = torch.randn_like(latents)
    bsz = latents.shape[0]
    # Sample a random timestep for each image
    timesteps = torch.randint(
        0,
        self.noise_scheduler.num_train_timesteps, (bsz, ),
        device=latents.device)
    timesteps = timesteps.long()
    # Add noise to the latents according to the noise magnitude at each timestep
    # (this is the forward diffusion process)
    noisy_latents = self.noise_scheduler.add_noise(latents, noise,
                                                   timesteps)
    input_ids = self.tokenize_caption(text).to(self.device)
    # Get the text embedding for conditioning
    with torch.no_grad():
        encoder_hidden_states = self.text_encoder(input_ids)[0]
    # Get the target for loss depending on the prediction type
    if self.noise_scheduler.config.prediction_type == 'epsilon':
        target = noise
    elif self.noise_scheduler.config.prediction_type == 'v_prediction':
        target = self.noise_scheduler.get_velocity(latents, noise,
                                                   timesteps)
    else:
        raise ValueError(
            f'Unknown prediction type {self.noise_scheduler.config.prediction_type}'
        )
    # Predict the noise residual and compute loss
    model_pred = self.unet(noisy_latents, timesteps,
                           encoder_hidden_states).sample
    if model_pred.shape[1] == 6:
        model_pred, _ = torch.chunk(model_pred, 2, dim=1)
    loss = F.mse_loss(model_pred.float(), target.float(), reduction='mean')
    output = {OutputKeys.LOSS: loss}
    return output