Does AMP work in TorchScript?

username: hughperkins
cooked: <p>Does AMP work in TorchScript?</p>

username: ptrblck
cooked: <p>If you could trace the model, all operations should be recorded (thus also the transformations) and amp should work. This would also mean that you don’t need to run the traced model in an autocast region anymore.</p>
<p>However, scripted models are currently still WIP.</p>

username: sidroopdaska
cooked: <p>Hey <a class="mention" href="/u/ptrblck">@ptrblck</a>, what do you mean by scripted models above? Would we need to trace the model within an autocast context block?</p>

username: ptrblck
cooked: <p>By “scripted model” I meant a model, which was scripted via <code>torch.jit.script</code>, while a traced model would be created via <code>torch.jit.trace</code>.<br>
As said before, tracing your model might work, if you don’t have any data-dependent code paths. I.e. if your model is static and tracing it works fine, you <em>should</em> be able to use amp. However, this is not tested properly and I would see it as an experimental workflow. Scripting a model with amp is not implemented yet.</p>

username: HaiNam
cooked: <p>Is the Script model currently usable with AMP? Thanks.</p>

username: ptrblck
cooked: <p>AMP should be working now in scripted models with the <code>nvfuser</code> JIT backend and by enabling it via <code>torch._C._jit_set_autocast_mode(True)</code>. You might need to install the latest nightly binary to get the current update of the <code>nvfuser</code> backend.</p>

username: HaiNam
cooked: <p>Thanks for the answer. But what is “latest nightly binary”? How can I install it. Thanks</p>

username: ptrblck
cooked: <p>You can pick it in the <a href="https://pytorch.org/get-started/locally/">install page</a> by selecting the “Preview (nightly)” option or my building from source.</p>

username: jayz
cooked: <aside class="quote no-group" data-username="ptrblck" data-post="7" data-topic="98079" data-full="true">
<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>AMP should be working now in scripted models with the <code>nvfuser</code> JIT backend and by enabling it via <code>torch._C._jit_set_autocast_mode(True)</code>. You might need to install the latest nightly binary to get the current update of the <code>nvfuser</code> backend.</p>
</blockquote>
</aside>
<p>Is this option only available in the nightly binaries?<br>
In my case:</p>
<pre data-code-wrap="python"><code class="lang-python">torch.__version__,torch.version.cuda,torch.backends.cudnn.version()
torch._C._jit_set_autocast_mode(True)
</code></pre>
<p>yields:</p>
<pre data-code-wrap="python"><code class="lang-python">---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
~\AppData\Local\Temp/ipykernel_17936/3336213180.py in &lt;module&gt;
      1 torch.__version__,torch.version.cuda,torch.backends.cudnn.version()
----&gt; 2 torch._C._jit_set_autocast_mode(True)

AttributeError: module 'torch._C' has no attribute '_jit_set_autocast_mode'
</code></pre>
<p>I get an error when using torch.jit.script with autocast, therefore I’m asking.<br>
Thanks!</p>