INNER CODE UNIT · Python
load_pretrained_lm_weights
lbcb-sci/RiNALMo · train_expression_level.py:53
def load_pretrained_lm_weights(self, pretrained_weights_path):
self.lm.load_state_dict(torch.load(pretrained_weights_path))
def forward(self, tokens):
x = self.lm(tokens)["representation"]
# Nullify padding token representations
pad_mask = tokens.eq(self.pad_idx)
x[pad_mask, :] = 0.0
pred = self.pred_head(x, pad_mask)
return pred
def _common_step(self, batch, batch_idx, log_prefix: str):
tokens, te_target = batch
preds = self(tokens)
loss = self.loss(preds, te_target)