def METHOD_NAME(self, set_to_none=False):
    if set_to_none:
        for group in self.param_groups:
            for p in group['params']:
                p.grad = None
    else:
        super(FusedAdam, self).METHOD_NAME()