format

open-spaced-repetition · Jan 25, 2025 · daaf5d1 · daaf5d1
1 parent b7066a3
commit daaf5d1
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/src/fsrs_optimizer/fsrs_optimizer.py b/src/fsrs_optimizer/fsrs_optimizer.py
@@ -392,7 +392,9 @@ def train(self, verbose: bool = True):
                 real_batch_size = seq_lens.shape[0]
                 outputs, _ = self.model(sequences)
                 stabilities = outputs[seq_lens - 1, torch.arange(real_batch_size), 0]
-                retentions = power_forgetting_curve(delta_ts, stabilities, -self.model.w[19])
+                retentions = power_forgetting_curve(
+                    delta_ts, stabilities, -self.model.w[19]
+                )
                 loss = (self.loss_fn(retentions, labels) * weights).sum()
                 penalty = torch.sum(
                     torch.square(self.model.w - self.init_w_tensor)