def METHOD_NAME(self, past_target: torch.Tensor) -> Distribution:
    # (batch_size, seq_len, target_dim) and (batch_size, seq_len, target_dim)
    scaled_target, target_scale = self.scaler(
        past_target,
        torch.ones_like(
            past_target
        ),  # TODO: pass the actual observed here
    )
    mlp_outputs = self.mlp(scaled_target)
    distr_args = self.distr_args_proj(mlp_outputs)
    return self.distr_output.distribution(
        distr_args, scale=target_scale.unsqueeze(1)
    )