INNER CODE UNIT · Python

forward

lbcb-sci/RiNALMo · train_ncrna_classification.py:55

    def forward(self, tokens):
        x = self.rinalmo(tokens)["representation"]
        x = x[:, 0] # take the CLS token, B x E
            
        pred = self.pred_head(x) # B x C
        return pred

    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)

View source record →