Torch: bugfix in diameter rescaling

ZIB-IOL · Oct 26, 2020 · 9afe7b9 · 9afe7b9
1 parent 875ea1b
commit 9afe7b9
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/frankwolfe/pytorch/optimizers.py b/frankwolfe/pytorch/optimizers.py
@@ -64,7 +64,7 @@ def step(self, constraints, closure=None):
 
                 if self.rescale == 'diameter':
                     # Rescale lr by diameter
-                    factor = 1. / constraints[idx].get_diameter(v.shape)
+                    factor = 1. / constraints[idx].get_diameter()
                 elif self.rescale == 'gradient':
                     # Rescale lr by gradient
                     factor = torch.norm(d_p, p=2) / torch.norm(p - v, p=2)