From f044114daecfecc49befb4928046ccd1963f81c1 Mon Sep 17 00:00:00 2001 From: Kevin Wu Date: Tue, 8 Nov 2022 21:56:18 -0800 Subject: [PATCH] Bugfix for model not learning --- bin/train_autoregressive.py | 2 +- foldingdiff/modelling.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/bin/train_autoregressive.py b/bin/train_autoregressive.py index 00e50f1..526fcc0 100644 --- a/bin/train_autoregressive.py +++ b/bin/train_autoregressive.py @@ -98,7 +98,7 @@ def train( gradient_clip: float = 1.0, batch_size: int = 256, lr: float = 5e-5, - l2_norm: float = 0.0, + l2_norm: float = 0.01, loss: modelling.LOSS_KEYS = "smooth_l1", min_epochs: Optional[int] = None, max_epochs: int = 10000, # 10000, set to 100 for debug diff --git a/foldingdiff/modelling.py b/foldingdiff/modelling.py index f8bf0b6..3e54a71 100644 --- a/foldingdiff/modelling.py +++ b/foldingdiff/modelling.py @@ -924,7 +924,7 @@ class BertForAutoregressive(BertForAutoregressiveBase, pl.LightningModule): assert preds.ndim == 3 # batch_size, seq_length, features # Get the loss terms l = self.loss( - preds[:, batch["causal_idx"]], + preds[torch.arange(batch["lengths"].shape[0]), batch["causal_idx"]], batch["causal_target"], ) return l