INNER CODE UNIT · Python
_validation_step
lbcb-sci/RiNALMo · train_ncrna_classification.py:76
def _validation_step(self, batch, batch_idx, split: str):
fam, tokens, labels = batch
pred = self(tokens)
loss = self.loss_fn(pred, labels)
self.val_step_outputs.append((torch.argmax(pred, dim=-1), labels))
log = {
f'{split}/loss': loss,
}
self.log_dict(log, sync_dist=True, add_dataloader_idx=False)
return loss
def _on_epoch_end(self, split: str):
preds = []