Fix optimizer import in gallery example.

cooper-org · May 6, 2022 · de1c6fe · de1c6fe
1 parent 5cd18c8
commit de1c6fe
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 205 deletions.
diff --git a/tutorials/Gaussian_mixture.py b/tutorials/Gaussian_mixture.py
diff --git a/tutorials/scripts/plot_gaussian_mixture.py b/tutorials/scripts/plot_gaussian_mixture.py
@@ -37,7 +37,6 @@
 from torch.nn.functional import binary_cross_entropy_with_logits as bce_loss
 
 import cooper
-from cooper.optim import SGD
 
 torch.manual_seed(0)
 np.random.seed(0)
@@ -160,10 +159,10 @@ def train(problem_name, inputs, targets, num_iters=5000, lr=1e-2, const_level=0.
     cmp = MixtureSeparation(is_constrained, use_proxy, const_level)
     formulation = cooper.LagrangianFormulation(cmp)
 
-    primal_optimizer = SGD(model.parameters(), lr=lr, momentum=0.7)
+    primal_optimizer = torch.optim.SGD(model.parameters(), lr=lr, momentum=0.7)
 
     if is_constrained:
-        dual_optimizer = cooper.optim.partial_optimizer(SGD, lr=lr, momentum=0.7)
+        dual_optimizer = cooper.optim.partial_optimizer(torch.optim.SGD, lr=lr, momentum=0.7)
     else:
         dual_optimizer = None