INNER CODE UNIT · Python

_common_step

lbcb-sci/RiNALMo · train_ncrna_classification.py:62

    def _common_step(self, batch, batch_idx, split):
        fam, tokens, labels = batch
        pred = self(tokens)

        loss = self.loss_fn(pred, labels)

        log = {
            f'{split}/loss': loss,
        }

        self.log_dict(log, sync_dist=True, add_dataloader_idx=False)

        return loss
    
    def _validation_step(self, batch, batch_idx, split: str):
        fam, tokens, labels = batch
        pred = self(tokens)

View source record →