diff options
| -rw-r--r-- | experiments/rain_ep_bias_train.py | 10 |
1 files changed, 6 insertions, 4 deletions
diff --git a/experiments/rain_ep_bias_train.py b/experiments/rain_ep_bias_train.py index eed5715..871d094 100644 --- a/experiments/rain_ep_bias_train.py +++ b/experiments/rain_ep_bias_train.py @@ -84,7 +84,6 @@ def main() -> None: from model.function.network import Network from model.hopfield.minimizer import FixedPointMinimizer from model.hopfield.network import ConvHopfieldEnergy28 - from training.monitor import Optimizer from training.sgd import AugmentedFunction, EquilibriumProp random.seed(args.seed) @@ -137,9 +136,12 @@ def main() -> None: inference_minimizer.mode = "asynchronous" inference_minimizer.num_iterations = args.inference_iterations learning_rates = [0.02] * len(energy.params()) - optimizer = Optimizer( - energy, cost, learning_rates, momentum=0.9, - weight_decay=3e-4) + parameter_groups = [ + {"params": parameter.state, "lr": learning_rate} + for parameter, learning_rate in zip(energy.params(), learning_rates) + ] + optimizer = torch.optim.SGD( + parameter_groups, lr=0.1, momentum=0.9, weight_decay=3e-4) metrics = [] start = time.time() |