Skip to content

Commit

Permalink
Update train_ema.py
Browse files Browse the repository at this point in the history
  • Loading branch information
smilesun authored Dec 5, 2024
1 parent 7d8d081 commit b296229
Showing 1 changed file with 1 addition and 1 deletion.
2 changes: 1 addition & 1 deletion domainlab/algos/trainers/train_ema.py
Original file line number Diff line number Diff line change
Expand Up @@ -57,7 +57,7 @@ def move_average(self, dict_data, epoch):
self._ma_iter += 1
return dict_return_ema_para_curr_iter

def after_epoch(self, epoch):
def after_epoch(self, epoch, flag_info=None):
torch_model = self.get_model()
dict_para = torch_model.state_dict() # only for trainable parameters
new_dict_para = self.move_average(dict_para, epoch)
Expand Down

0 comments on commit b296229

Please sign in to comment.