# Buggy code
def assign_average_vars(self, var_list):
    """Assign variables in var_list with their respective averages.

    Args:
        var_list: List of model variables to be assigned to their average.

    Returns:
        assign_op: The op corresponding to the assignment operation of
        variables to their average.

    Example:
    ```python
    model = tf.Sequential([...])
    opt = tfa.optimizers.SWA(
            tf.keras.optimizers.SGD(lr=2.0), 100, 10)
    model.compile(opt, ...)
    model.fit(x, y, ...)

    # Update the weights to their mean before saving
    opt.assign_average_vars(model.variables)

    model.save('model.h5')
    ```
    """
    assign_op = tf.group(
        [
            var.assign(self.get_slot(var, "average"), use_locking=self._use_locking)
            for var in var_list
            if var.trainable
        ]
    )
    return assign_op


# Traceback (AttributeError)
Traceback (most recent call last):
File "/usr/lib/python3.6/runpy.py", line 193, in _run_module_as_main
"__main__", mod_spec)
File "/usr/lib/python3.6/runpy.py", line 85, in _run_code
exec(code, run_globals)
File "/root/.vscode-server/extensions/ms-python.python-2020.11.371526539/pythonFiles/lib/python/debugpy/__main__.py", line 45, in <module>
cli.main()
File "/root/.vscode-server/extensions/ms-python.python-2020.11.371526539/pythonFiles/lib/python/debugpy/../debugpy/server/cli.py", line 430, in main
run()
File "/root/.vscode-server/extensions/ms-python.python-2020.11.371526539/pythonFiles/lib/python/debugpy/../debugpy/server/cli.py", line 267, in run_file
runpy.run_path(options.target, run_name=compat.force_str("__main__"))
File "/usr/lib/python3.6/runpy.py", line 263, in run_path
pkg_name=pkg_name, script_name=fname)
File "/usr/lib/python3.6/runpy.py", line 96, in _run_module_code
mod_name, mod_spec, pkg_name, script_name)
File "/usr/lib/python3.6/runpy.py", line 85, in _run_code
exec(code, run_globals)
File "/usr/local/srv/automl/efficientdet/nlp.py", line 65, in <module>
tfa.callbacks.AverageModelCheckpoint(update_weights=True, filepath=os.path.join(MODEL_DIR, 'ckpt_moving_average'))])
File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/training.py", line 108, in _method_wrapper
return method(self, *args, **kwargs)
File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/engine/training.py", line 1137, in fit
callbacks.on_epoch_end(epoch, epoch_logs)
File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/callbacks.py", line 412, in on_epoch_end
callback.on_epoch_end(epoch, logs)
File "/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/callbacks.py", line 1249, in on_epoch_end
self._save_model(epoch=epoch, logs=logs)
File "/usr/local/lib/python3.6/dist-packages/tensorflow_addons/callbacks/average_model_checkpoint.py", line 76, in _save_model
self.model.optimizer.assign_average_vars(self.model.trainable_weights)
File "/usr/local/lib/python3.6/dist-packages/tensorflow_addons/optimizers/average_wrapper.py", line 125, in assign_average_vars
for var in var_list
File "/usr/local/lib/python3.6/dist-packages/tensorflow_addons/optimizers/average_wrapper.py", line 126, in <listcomp>
if var.trainable
AttributeError: 'TrackableWeightHandler' object has no attribute 'trainable'

# Fixed code
def assign_average_vars(self, var_list):
    """Assign variables in var_list with their respective averages.

    Args:
        var_list: List of model variables to be assigned to their average.

    Returns:
        assign_op: The op corresponding to the assignment operation of
        variables to their average.

    Example:
    ```python
    model = tf.Sequential([...])
    opt = tfa.optimizers.SWA(
            tf.keras.optimizers.SGD(lr=2.0), 100, 10)
    model.compile(opt, ...)
    model.fit(x, y, ...)

    # Update the weights to their mean before saving
    opt.assign_average_vars(model.variables)

    model.save('model.h5')
    ```
    """
    assign_ops = []
    for var in var_list:
        try:
            assign_ops.append(
                var.assign(self.get_slot(var, "average"), use_locking=self._use_locking)
            )
        except Exception as e:
            warnings.warn("Unable to assign average slot to {} : {}".format(var, e))
    return tf.group(assign_ops)