NameError: name ‘custom_fwd’ is not defined

Sureerat_Reaungamorn: <p>Hi, I have a custom torch.autograd.Function function and try to use mix precision training but I got error  below. I use pytorch 1.7</p>
<p><code>NameError: name ‘custom_fwd’ is not defined</code></p>
<p>Here is the example code</p>
<pre><code class="lang-auto">class MyFloat32Func(torch.autograd.Function):
    @staticmethod
    @custom_fwd(cast_inputs=torch.float32)
    def forward(ctx, input):
        ctx.save_for_backward(input)
        pass
        return fwd_output
    @staticmethod
    @custom_bwd
    def backward(ctx, grad):
        pass
    
func=MyFloat32Func.apply
func(2)
</code></pre>
<p>And I got error</p>
<pre><code class="lang-auto">---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
&lt;ipython-input-3-744176e249fa&gt; in &lt;module&gt;
----&gt; 1 class MyFloat32Func(torch.autograd.Function):
      2     @staticmethod
      3     @custom_fwd(cast_inputs=torch.float32)
      4     def forward(ctx, input):
      5         ctx.save_for_backward(input)

&lt;ipython-input-3-744176e249fa&gt; in MyFloat32Func()
      1 class MyFloat32Func(torch.autograd.Function):
      2     @staticmethod
----&gt; 3     @custom_fwd(cast_inputs=torch.float32)
      4     def forward(ctx, input):
      5         ctx.save_for_backward(input)

NameError: name 'custom_fwd' is not defined
</code></pre>
<p>If I call the following print</p>
<pre><code class="lang-auto">print(torch.cuda.amp.custom_fwd)
print(torch.cuda.amp.custom_fwd(cast_inputs=torch.float32))
</code></pre>
<p>I get</p>
<pre><code class="lang-auto">&lt;function custom_fwd at 0x7f9b0310cef0&gt;
functools.partial(&lt;function custom_fwd at 0x7f9b0310cef0&gt;, cast_inputs=torch.float32)
</code></pre>

ptrblck: <p>You would have to either import <code>custom_fwd</code> from <code>torch.cuda.amp</code> or specify the full namespace as <code>torch.cuda.amp.custom_fwd</code>.</p>

Sureerat_Reaungamorn: <p>Thank you very much! <code>@torch.cuda.amp.custom_fwd</code> works!!</p>
<pre><code class="lang-auto">class MyFloat32Func(torch.autograd.Function):
    @staticmethod
    @torch.cuda.amp.custom_fwd(cast_inputs=torch.float32)
    def forward(ctx, input):
        ctx.save_for_backward(input)
        return torch.tensor(input,dtype=torch.float32)
    @staticmethod
    @torch.cuda.amp.custom_fwd
    def backward(ctx, grad):
        return 1
    
func=MyFloat32Func.apply
func(2)
</code></pre>
<p>Just wondering, why import <code>custom_fwd</code> from <code>torch.cuda.amp</code> does not work…</p>
<pre><code class="lang-auto">from torch.cuda.amp import custom_fwd


class MyFloat32Func(torch.autograd.Function):
    @staticmethod
    @custom_fwd(cast_inputs=torch.float32)
    def forward(ctx, input):
        ctx.save_for_backward(input)
        return torch.tensor(input,dtype=torch.float32)
    @staticmethod
    @custom_bwd
    def backward(ctx, grad):
        return 1
    
func=MyFloat32Func.apply
func(2)
</code></pre>
<p>Here is the error</p>
<pre><code class="lang-auto">---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
&lt;ipython-input-12-d1461ca0ff6e&gt; in &lt;module&gt;
      2 
      3 
----&gt; 4 class MyFloat32Func(torch.autograd.Function):
      5     @staticmethod
      6     @custom_fwd(cast_inputs=torch.float32)

&lt;ipython-input-12-d1461ca0ff6e&gt; in MyFloat32Func()
      9         return torch.tensor(input,dtype=torch.float32)
     10     @staticmethod
---&gt; 11     @custom_bwd
     12     def backward(ctx, grad):
     13         return 1

NameError: name 'custom_bwd' is not defined

</code></pre>

ptrblck: <p>You are only importing <code>custom_fwd</code> in the second approach, while you would need to import <code>custom_fwd</code> and <code>custom_bwd</code>.</p>
<p>Also, in your first example you are using <code>custom_fwd</code> twice, so you might want to change it as well.</p>

Sureerat_Reaungamorn: <aside class="quote no-group" data-username="ptrblck" data-post="4" data-topic="122240">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>custom_bwd</p>
</blockquote>
</aside>
<p>Thank you very much!! It works fine now with <code>custom_bwd</code> above <code>backward</code> function!! <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>