LsrModelOutput

class LsrModelOutput(prediction: torch.FloatTensor, loss: Optional[torch.FloatTensor] = None)

Output type of LsrModel

Parameters

• prediction (torch.FloatTensor of shape (batch_size, max_h_t_count, num_relations)) – Prediction scores for all head to tail entity combinations from the final layer. Note that the sigmoid function has not been applied at this point.

• loss (torch.FloatTensor of shape (1,), optional, returned when labels is provided) – Loss on relation prediction task.