Class autograd.Function in module cause autoreload fail in jupyter lab

username: hu_jack
cooked: <p>I define a <strong>foo.py</strong> file and then creat a jupyter lab notebook that will import the class CNN in this file.<br>
In this <strong>foo.py</strong> file. I first define the <code>class CNN</code> then the <code>class fun</code></p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import torch.nn.functional as F

class CNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1,32,3,1,1) 
        print('pytorch')
        print('ok')
        
    
class fun(torch.autograd.Function):

    @staticmethod
    def forward(ctx, input):
        ctx.save_for_backward(input)
        return input.gt(thresh).float()

    @staticmethod
    def backward(ctx, grad_output):
        input, = ctx.saved_tensors
        grad_input = grad_output.clone()
        temp = abs(input - thresh) &lt; lens
        return grad_input * temp.float()
</code></pre>
<p>In the jupyter lab, I enable the autoreload</p>
<pre><code class="lang-auto">%load_ext autoreload
%autoreload 2
</code></pre>
<p>And the autoreload magic will work. <strong>However,</strong> if I change the sequence of class defined in <strong>foo.py</strong>,<br>
which means I first define <code>class fun</code> then define <code>class CNN</code>. The autoreload will not work in my jupyter lab notebook.</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import torch.nn.functional as F

class fun(torch.autograd.Function):

    @staticmethod
    def forward(ctx, input):
        ctx.save_for_backward(input)
        return input.gt(thresh).float()

    @staticmethod
    def backward(ctx, grad_output):
        input, = ctx.saved_tensors
        grad_input = grad_output.clone()
        temp = abs(input - thresh) &lt; lens
        return grad_input * temp.float()
    
    
class CNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1,32,3,1,1) 
        print('pytorch')
        print('ok')
</code></pre>
<p>In this case, If I modify the <code>__init__</code>  in <code>class CNN</code>, the autoreload didn’t work.<br>
Anyone can answer why ？</p>

username: albanD
cooked: <p>Hi,</p>
<p>Quite weird indeed.<br>
Does it only happen in you define a Function subclass? What if you just define any class:</p>
<pre><code class="lang-auto">class Dummy():
    pass
</code></pre>

username: hu_jack
cooked: <p>I found that:<br>
if I define as <code>class fun()</code> instead of<code> class fun(torch.autograd.Function):</code><br>
the autoreload succeed.</p>
<p>Also I try following , the autoreload works for both sequences.</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import torch.nn.functional as F

class fun(torch.autograd.Function):

    @staticmethod
    def forward(ctx, input):
        ctx.save_for_backward(input)
        return input.gt(thresh).float()

    @staticmethod
    def backward(ctx, grad_output):
        input, = ctx.saved_tensors
        grad_input = grad_output.clone()
        temp = abs(input - thresh) &lt; lens
        return grad_input * temp.float()

class Dummy():
    print("torch")
</code></pre>
<p>`</p>

username: albanD
cooked: <p>Hi,</p>
<p>I don’t think pytorch supports reloading the module. So this is most likely going to fail in unexpacted ways if you try to <code>reload(torch)</code> no?</p>

username: hu_jack
cooked: <p>Really thanks for your help. I don’t understand, where to<code> reload(torch)</code>? and why need <code>reload(torch)</code>, because I only modify the module that I write and the torch module is untouched.</p>

username: albanD
cooked: <p>I am not sure how <code>autoreload</code> works but would it be reloading torch when it reload your module?</p>

username: hu_jack
cooked: <p>Thank you, may be I should post this to ipython in github.<br>
the Reload all modules (except those excluded by  <code>%aimport</code> ) every time before executing the Python code typed.<br>
This is what the official docs say</p>

username: albanD
cooked: <p>I know that using <code>importlib.reload(torch)</code> will not behave properly (and usually ends up segfaulting), we do have an issue open to track this out. But this will be quite hard to fix.</p>