def forward(self, x, timesteps):
    '\n        Apply the model to an input batch.\n\n        :param x: an [N x C x ...] Tensor of inputs.\n        :param timesteps: a 1-D batch of timesteps.\n        :return: an [N x K] Tensor of outputs.\n        '
    emb = self.time_embed(timestep_embedding(timesteps, self.model_channels))
    results = []
    h = x.type(self.dtype)
    for module in self.input_blocks:
        h = module(h, emb)
    h = self.middle_block(h, emb).type(self.dtype)
    image_features = self.out(h)
    image_features = F.normalize(image_features, dim=(- 1))
    return image_features
Apply the model to an input batch.

:param x: an [N x C x ...] Tensor of inputs.
:param timesteps: a 1-D batch of timesteps.
:return: an [N x K] Tensor of outputs.
def forward(self, x, timesteps):
    '\n        Apply the model to an input batch.\n\n        :param x: an [N x C x ...] Tensor of inputs.\n        :param timesteps: a 1-D batch of timesteps.\n        :return: an [N x K] Tensor of outputs.\n        '
    emb = self.time_embed(timestep_embedding(timesteps, self.model_channels))
    results = []
    h = x.type(self.dtype)
    for module in self.input_blocks:
        h = module(h, emb)
    h = self.middle_block(h, emb).type(self.dtype)
    image_features = self.out(h)
    image_features = F.normalize(image_features, dim=(- 1))
    return image_features