suppress unnecessary warnings (#2333)

tensorflow · Oct 9, 2021 · 20eb284 · 20eb284
1 parent 8e73b9d
commit 20eb284
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tensorflow_addons/callbacks/average_model_checkpoint.py b/tensorflow_addons/callbacks/average_model_checkpoint.py
@@ -82,13 +82,13 @@ def _save_model(self, *args, **kwargs):
  assert isinstance(optimizer, AveragedOptimizerWrapper)
 
  if self.update_weights:
- optimizer.assign_average_vars(self.model.variables)
+ optimizer.assign_average_vars(self.model.trainable_weights)
  return super()._save_model(*args, **kwargs)
  else:
  # Note: `model.get_weights()` gives us the weights (non-ref)
  # whereas `model.variables` returns references to the variables.
  non_avg_weights = self.model.get_weights()
- optimizer.assign_average_vars(self.model.variables)
+ optimizer.assign_average_vars(self.model.trainable_weights)
  # result is currently None, since `super._save_model` doesn't
  # return anything, but this may change in the future.
  result = super()._save_model(*args, **kwargs)