fix: Correct learning rate update in optimizer, reconstruction, surrogate

modules/optimizer.py

@@ -88,7 +88,8 @@ class Optimizer(object):
         using the reconstruction model loss
         '''
         # set the optimizer
-        self.optimizer.lr = lr
+        for param_group in self.optimizer.param_groups:
+            param_group['lr'] = lr
 
         self.surrogate_model.eval()
         self.reconstruction_model.eval()

modules/reconstruction.py

@@ -52,7 +52,8 @@ class Reconstruction(torch.nn.Module):
 
         train_loader = DataLoader(dataset,  batch_size=batch_size, shuffle=True)
         # set the optimizer
-        self.optimizer.lr = lr
+        for param_group in self.optimizer.param_groups:
+            param_group['lr'] = lr
         self.to(self.device)
         self.train()
         for epoch in range(n_epochs):

modules/surrogate.py

@@ -158,7 +158,8 @@ class Surrogate(torch.nn.Module):
         # train the surrogate model
         train_loader = DataLoader(surrogate_dataset,  batch_size=batch_size, shuffle=True)
         # set the optimizer
-        self.optimizer.lr = lr
+        for param_group in self.optimizer.param_groups:
+            param_group['lr'] = lr
         self.to(self.device)
         
         self.train()