How to update parameter with moving average method

[hustzxd]

I want to add a parameter in my custom layer, and I want to update it by moving average. So far, I have implemented a simple version, but it can't perform hybridize. Hope some body's help.

class CustomLayer(HybridBlock):
    def __init__(self, prefix=None, params=None):
        super(CustomLayer, self).__init__(prefix=prefix, params=params)
        init_scale = mx.init.Constant([32])
        with self.name_scope():
            self.scale = self.params.get('scale',
                                         shape=[1],
                                         init=init_scale,
                                         allow_deferred_init=False,
                                         differentiable=False)

    def hybrid_forward(self, F, x, scale=None):
        if autograd.is_training():
            with autograd.pause():
                scale = scale * 0.9 + new_scale * 0.1
                self.scale.set_data(scale)
        y = x
        return y

[zhreshold]

I think you don't need to call set_data explicitly in hybrid_forward, which is not recommended way in the computation graph.

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.