swa links

sparks-baird · Feb 22, 2022 · 13b4ac7 · 13b4ac7
1 parent 9369c34
commit 13b4ac7
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/crabnet/crabnet_.py b/crabnet/crabnet_.py
@@ -379,6 +379,8 @@ def _train(self):
                     )
                     pred_v = np.nan_to_num(pred_v)
                 mae_v = mean_absolute_error(true_v, pred_v)
+                # https://github.com/pytorch/contrib/blob/master/torchcontrib/optim/swa.py
+                # https://pytorch.org/blog/pytorch-1.6-now-includes-stochastic-weight-averaging/
                 self.optimizer.update_swa(mae_v)
                 minima.append(self.optimizer.minimum_found)