Thank you for your responses.
@michaelklachko: I need the flexibility of being agnostic to the actual module to which DropConnect is applied. That is, it should work for any torch.nn module that has a weight parameter (nn.Linear, nn.Conv2D etc.). Your simpler method is in my case not really simpler as it requires me to put “module code” into the training loop which I would like to avoid to keep the usage generic.
@rwightman: Your suggestions assume that I have a single scalar drop_prob value, but in my case I need to control this parameter for every weight individually. That is why I am using the full tensor. I register it as a buffer to make sure that PyTorch properly moves it to the correct device when to() is called.