In this previous thread 11, an implementation of DropConnect is discussed. The recommended method 12 assumes that the dropout probability remains constant. I am now trying to extend this to the case where every parameter has its own dropout probability and the dropout probabilities are updated during training according to some learning rule.
In my attempt so far, I register an additional buffer for each module that contains the dropout probabilities. However, I have trouble then updating the probabilities during training. To update the probabilities I use the following way of accessing the buffer:
old_probs = getattr(module, param_name + "_prob")
new_probs = learning_rule(old_probs)
setattr(module, param_name + "_prob", new_probs)

The updates don’t seem to work as intended. Can someone tell me whether this is a reasonable approach, or if I am missing something?
Here is the slightly modified _weight_drop() function with the additional buffer:
def _weight_drop(module, weights, dropout):
    """
    Helper for `WeightDrop`.
    """

    for name_w in weights:
        w = getattr(module, name_w)
        del module._parameters[name_w]
        module.register_parameter(name_w + '_raw', Parameter(w))
        module.register_buffer(name_w + '_prob', torch.full_like(w, dropout))

    original_module_forward = module.forward

    def forward(*args, **kwargs):
        for name_w in weights:
            raw_w = getattr(module, name_w + '_raw')
            prob_w = getattr(module, name_w + '_prob')
            w = torch.bernoulli(1.0 - prob_w) * raw_w
            setattr(module, name_w, w)

        return original_module_forward(*args, **kwargs)

    setattr(module, 'forward', forward)