Dropconnect layer implementation.

[swaroop1904]

I am trying to code dropconnect for Conv2D and transposeconv2D layer by following the implementation for dropconnect lstm and dropconnect gru in this repo. Below is my implementation of the code.

def _weight_drop(module, weights, dropout):
    for name_w in weights:
        w = getattr(module, name_w)
        del module._parameters[name_w]
        module.register_parameter(name_w + '_raw', Parameter(w))
    original_module_forward = module.forward

    def forward(*args, **kwargs):
        for name_w in weights:
            raw_w = getattr(module, name_w + '_raw')
            w = torch.nn.functional.dropout(raw_w, p=dropout, training=module.training)
            setattr(module, name_w, w)
        return original_module_forward(*args, **kwargs)
    setattr(module, 'forward', forward)
        
class WeightDropConv2d(torch.nn.Conv2d):
    def __init__(self, *args, weight_dropout=0.0, **kwargs):
        super().__init__(*args, **kwargs)
        weights = ['weight']
        _weight_drop(self, weights, weight_dropout)
        
class WeightDropConvTranspose2d(torch.nn.ConvTranspose2d):
    def __init__(self, *args, weight_dropout=0.0, **kwargs):
        super().__init__(*args, **kwargs)
        weights = ['weight']
        _weight_drop(self, weights, weight_dropout)

torch.version.cuda: 1.1.0
torch.version: 9.0.176

I get the following error in the 2nd epoch:

Traceback (most recent call last):
  File "dropconnect.py", line 110, in <module>
    out = model(image)
  File "/home/sbhand2s/.local/lib/python3.6/site-packages/torch/nn/modules/module.py", line 493, in __call__
    result = self.forward(*input, **kwargs)
  File "dropconnect.py", line 73, in forward
    out = self.c1(x)
  File "/home/sbhand2s/.local/lib/python3.6/site-packages/torch/nn/modules/module.py", line 493, in __call__
    result = self.forward(*input, **kwargs)
  File "dropconnect.py", line 34, in forward
    setattr(module, name_w, w)
  File "/home/sbhand2s/.local/lib/python3.6/site-packages/torch/nn/modules/module.py", line 558, in __setattr__
    .format(torch.typename(value), name))
TypeError: cannot assign 'torch.cuda.FloatTensor' as parameter 'weight' (torch.nn.Parameter or None expected)

This error occurs in the second epoch when I switch from .eval() to .train(). This error does not occur if I don't call .eval()

Any suggestions on why this error is occurring or how to implement dropconnect in a better manner?