kosmos.ml.models.model

Classes

class Model(input_dim: int, output_dim: int)

Bases: torch.nn.Module, abc.ABC

Base class for models.

Initialize the model.

Parameters:

• input_dim (int) – Model input dimension.

• output_dim (int) – Model output dimension.

Methods

forward(x: torch.Tensor) → torch.Tensor

Perform a forward pass.

Parameters:

x (torch.Tensor) – Input tensor.

Returns:

Output tensor.

Return type:

torch.Tensor