INNER CODE UNIT · Python
load_pretrained_rinalmo_weights
lbcb-sci/RiNALMo · train_ncrna_classification.py:52
def load_pretrained_rinalmo_weights(self, pretrained_weights_path):
self.rinalmo.load_state_dict(torch.load(pretrained_weights_path))
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,