How can I save my model when pickling fails?

username: HarrisonWilde
cooked: <p>How can I make my modules compliant with <code>torch.save</code> in order to avoid this error that I see when I try to save my model after training?</p>
<pre><code class="lang-auto">Traceback (most recent call last):
  File "main.py", line 351, in &lt;module&gt;
    main()
  File "main.py", line 225, in main
    synth_data, log_model_iw = model_map[args.model_class](**params)
  File "/Users/harrisonwilde/Library/Mobile Documents/com~apple~CloudDocs/PhD/Holmes/WeightedDP/models/CUSTOM/dpgan.py", line 250, in DPGAN_runner
    torch.save(gan, save)
  File "/usr/local/Caskroom/miniconda/base/envs/dp/lib/python3.8/site-packages/torch/serialization.py", line 372, in save
    _save(obj, opened_zipfile, pickle_module, pickle_protocol)
  File "/usr/local/Caskroom/miniconda/base/envs/dp/lib/python3.8/site-packages/torch/serialization.py", line 476, in _save
    pickler.dump(obj)
TypeError: cannot pickle 'torch._C.Generator' object
</code></pre>
<p>Forgive the vagueness of the question, I am pretty new to PyTorch but cannot find anything about this error online. The Generator module I have written is as follows:</p>
<pre><code class="lang-auto">class Generator(nn.Module):
    def __init__(self, latent_dim, output_dim, binary=True):
        super(Generator, self).__init__()

        def block(in_, out, Activation):
            return nn.Sequential(
                nn.Linear(in_, out, bias=False),
                nn.LayerNorm(out),
                Activation(),
            )

        self.layer_0 = block(latent_dim, latent_dim, nn.Tanh if binary else lambda: nn.LeakyReLU(0.2))
        self.layer_1 = block(latent_dim, latent_dim, nn.Tanh if binary else lambda: nn.LeakyReLU(0.2))
        self.layer_2 = block(latent_dim, output_dim, nn.Tanh if binary else lambda: nn.LeakyReLU(0.2))

    def forward(self, noise):
        noise = self.layer_0(noise) + noise
        noise = self.layer_1(noise) + noise
        noise = self.layer_2(noise)
        return noise
</code></pre>
<p>Would appreciate any pointers, thanks!</p>

username: mwalston
cooked: <p>I am also somewhat new to this, but I believe your error is because your <code>block(in_, out, Activation)</code> function is defined as a nested function inside <code>init()</code> and according to the discussion at <a href="https://stackoverflow.com/questions/12019961/python-pickling-nested-functions" class="inline-onebox" rel="noopener nofollow ugc">Python: pickling nested functions - Stack Overflow</a> python can not pickle nested functions.</p>
<p>A solution would be to move <code>def block</code> out into the class scope like this:</p>
<pre><code class="lang-auto">class Generator(nn.Module):
    def __init__(self, latent_dim, output_dim, binary=True):
        super(Generator, self).__init__()

        self.layer_0 = block(latent_dim, latent_dim, nn.Tanh if binary else lambda: nn.LeakyReLU(0.2))
        self.layer_1 = block(latent_dim, latent_dim, nn.Tanh if binary else lambda: nn.LeakyReLU(0.2))
        self.layer_2 = block(latent_dim, output_dim, nn.Tanh if binary else lambda: nn.LeakyReLU(0.2))

    def forward(self, noise):
        noise = self.layer_0(noise) + noise
        noise = self.layer_1(noise) + noise
        noise = self.layer_2(noise)
        return noise

    def block(in_, out, Activation):
        return nn.Sequential(
            nn.Linear(in_, out, bias=False),
            nn.LayerNorm(out),
            Activation(),
        )</code></pre>

username: HarrisonWilde
cooked: <p>Hi, thank you, I think you are right, I rearranged my code and have tried the block function in a few places but this doesn’t seem to fix it</p>

username: acxz
cooked: <p>Do you have a self.torch_seed call? That could do it. Try to set the seed outside of the object you are trying to pickle.</p>