INNER CODE UNIT · Python

_common_step

lbcb-sci/RiNALMo · train_expression_level.py:66

    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)
        self.r2_metric.update(preds, te_target)
        self.spearman_metric.update(preds, te_target)

        log = {
            f'{log_prefix}/loss': loss,
            f'{log_prefix}/mse': mse,
            f'{log_prefix}/mae': mae,
        }
        self.log_dict(log, sync_dist=True)

View source record →