mtenn.model.LigandOnlyModel

class mtenn.model.LigandOnlyModel(model, readout=None, fix_device=False)

Bases: Model

A ligand-only version of the Model. In this case, the representation block will hold the entire model, while the strategy block will simply be set as an Identity module.

Methods

__init__(model[, readout, fix_device])	Build a LigandOnlyModel.
forward(rep)	Handles all the logic detailed in the docs page.

__init__(model, readout=None, fix_device=False)

Build a LigandOnlyModel.

Parameters:

• model – This can be any kind of model that will go from a single input representation to a prediction (eg a GAT <mtenn.conversion_utils.gat.GAT instance)

• fix_device (bool, default=False) – If True, make sure the input is on the same device as the model, copying over as necessary.

forward(rep)

Handles all the logic detailed in the docs page.

Parameters:

rep – Whatever input representation the unerlying model takes

Returns:

• torch.Tensor – Final model prediction. If the model has a readout, this value will have the readout applied

• list[torch.Tensor] – A list containing only the pre-readout model prediction. This value is returned mainly to align the signatures for the single- and multi-pose models