Multiprocessing pickling error

Hi,

I have a SpatialDoupout class like this:

import os
import numpy as np
from torch import nn
import torch
from torch.autograd import Variable
from torch.nn._functions.dropout import Dropout
from itertools import repeat
from torch.nn.modules.dropout import _DropoutNd

class _SpatialDropout(Dropout):
    def __init__(self):
        super(_SpatialDropout,self).__init__()
    @staticmethod
    def symbolic(g, input, p=0.5, train=False, inplace=False):
        # See Note [Export inplace]
        # NB: In inference mode, FeatureDropout is exported as an identity op.
        from torch.onnx.symbolic import _unimplemented
        if train:
            return _unimplemented("FeatureDropout", "training mode")
        return input

    @staticmethod
    def _make_noise(input):
        print('input:')
        print(input)
        return input.new().resize_(input.size()[0],
                                    *repeat(1, input.dim() - 2),
                                    input.size()[2])

class SpatialDropout(_DropoutNd):
    def __init__(self,p):
        super(SpatialDropout,self).__init__()
        self.p = p
        self.function = _SpatialDropout()

    def forward(self, input):
        return self.function.apply(input, self.p, self.training, self.inplace)

When I use multiprocessing, got an error:

TypeError: can’t pickle _SpatialDropout objects

How can I fix it?

Thanks