I using Pytorch’s swa_utils that internally calls deepcopy.
However, I get the following error:
> Only Tensors created explicitly by the user (graph leaves) 
support the deepcopy protocol at the moment

This seems to be caused by the weight dropout scheme I am using, where the original weight matrix is moved into _raw. The code for weight dropout is here:
import torch
from torch.nn import Parameter
import torch.nn.functional as F 

class WeightDrop(object):
    def __init__(self, name, dropout):
        self.name = name
        self.dropout = dropout

    def compute_weight(self, module):
        return F.dropout(
            getattr(module, self.name + "_raw"),
            p        = self.dropout,
            training = module.training,
            inplace  = False
        )

    @staticmethod
    def apply(module, name, dropout):
        for k, hook in module._forward_pre_hooks.items():
            if isinstance(hook, WeightDrop) and hook.name == name:
                raise RuntimeError(f"Cannot register two weight_dropout hooks with name '{name}'")
        fn = WeightDrop(name, dropout)
        weight = getattr(module, name)

        del module._parameters[name]

        ## creating _raw parameter
        module.register_parameter(name + "_raw", Parameter(weight.data))
        setattr(module, name, fn.compute_weight(module))

        module.register_forward_pre_hook(fn)
        return fn

    def remove(self, module):
        weight = module._parameters[self.name + "_raw"]
        delattr(module, self.name)
        del module._parameters[self.name + "_raw"]
        module.register_parameter(self.name, Parameter(weight.data))

    def __call__(self, module, inputs):
        if self.name in module._parameters:
            del module._parameters[self.name]
        setattr(module, self.name, self.compute_weight(module))


def weight_drop(module, name, dropout):
    WeightDrop.apply(module, name, dropout)
    return module

def apply_weight_drop(module, name, dropout):
    wdrop = weight_drop(module, name, dropout)
    fp = module.flatten_parameters

    def decorator(*args, **kwargs):
        device = getattr(module, name + "_raw").device
        setattr(module, name, getattr(module, name).to(device))
        return fp(*args, **kwargs)

    module.flatten_parameters = decorator
    return wdrop

Here is the code for applying deep copy on a GRU:
import copy
gru = torch.nn.GRU(10, 10)
gru_wd = apply_weight_drop(gru, "weight_hh_l0", 0.2)
gru_wd_copy = copy.deepcopy(gru_wd)
> RuntimeError: Only Tensors created explicitly by the user (graph leaves) support the deepcopy protocol at the moment

Any ideas how to fix this error?
Thank you!