INNER CODE UNIT · Python

forward

lbcb-sci/RiNALMo · train_expression_level.py:56

    def forward(self, tokens):
        x = self.lm(tokens)["representation"]

        # Nullify padding token representations
        pad_mask = tokens.eq(self.pad_idx)
        x[pad_mask, :] = 0.0

        pred = self.pred_head(x, pad_mask)
        return pred

    def _common_step(self, batch, batch_idx, log_prefix: str):
        tokens, te_target = batch
        preds = self(tokens)

        loss = self.loss(preds, te_target)

        mse = F.mse_loss(preds, te_target)
        mae = F.l1_loss(preds, te_target)

View source record →