mxtaltools.models.task_models.embedding_regression_models

class mxtaltools.models.task_models.embedding_regression_models.EquivariantEmbeddingRegressor(seed, config, num_targets: int = 1, conditions_dim: int = 0, prediction_type: str = 'scalar')[source]

Bases: BaseGraphModel

single property prediction head for pretrained embeddings

forward(x: Tensor, v: Tensor) → tuple[Tensor, Tensor][source]: no need to do standardization, inputs are raw outputs from autoencoder model

class mxtaltools.models.task_models.embedding_regression_models.InvariantEmbeddingRegressor(seed, config, num_targets: int = 1, conditions_dim: int = 0, target_standardization_tensor: Tensor | None = None)[source]

Bases: BaseGraphModel

single property prediction head for pretrained embeddings

forward(x: Tensor) → Tensor[source]: no need to do standardization, inputs are raw outputs from autoencoder model