remove update_ema_freq

agent/pretrain/train_agent.py

@@ -83,7 +83,6 @@ class PreTrainAgent:
         # Training params
         self.n_epochs = cfg.train.n_epochs
         self.batch_size = cfg.train.batch_size
-        self.update_ema_freq = cfg.train.update_ema_freq
         self.epoch_start_ema = cfg.train.epoch_start_ema
         self.val_freq = cfg.train.get("val_freq", 100)
 

agent/pretrain/train_diffusion_agent.py

@@ -38,8 +38,7 @@ class TrainDiffusionAgent(PreTrainAgent):
                 self.optimizer.zero_grad()
 
                 # update ema
-                if self.epoch % self.update_ema_freq == 0:
-                    self.step_ema()
+                self.step_ema()
             loss_train = np.mean(loss_train_epoch)
 
             # validate

agent/pretrain/train_gaussian_agent.py

@@ -46,8 +46,7 @@ class TrainGaussianAgent(PreTrainAgent):
                 self.optimizer.zero_grad()
 
                 # update ema
-                if self.epoch % self.update_ema_freq == 0:
-                    self.step_ema()
+                self.step_ema()
             loss_train = np.mean(loss_train_epoch)
             ent_train = np.mean(ent_train_epoch)