1.13/amp.html



<!DOCTYPE html>
<!--[if IE 8]><html class="no-js lt-ie9" lang="en" > <![endif]-->
<!--[if gt IE 8]><!--> <html class="no-js" lang="en" > <!--<![endif]-->
<head>
  <meta name="robots" content="noindex">
  <meta charset="utf-8">
  <meta name="generator" content="Docutils 0.18.1: http://docutils.sourceforge.net/" />

  <meta name="viewport" content="width=device-width, initial-scale=1.0">
  
  <title>Automatic Mixed Precision package - torch.amp &mdash; PyTorch 1.13 documentation</title>
  

    <link rel="canonical" href="https://pytorch.org/docs/stable/amp.html"/>
  

  <link rel="stylesheet" href="_static/css/theme.css" type="text/css" />
  <!-- <link rel="stylesheet" href="_static/pygments.css" type="text/css" /> -->
  <link rel="stylesheet" href="_static/pygments.css" type="text/css" />
  <link rel="stylesheet" href="_static/css/theme.css" type="text/css" />
  <link rel="stylesheet" href="_static/copybutton.css" type="text/css" />
  <link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.10.0-beta/dist/katex.min.css" type="text/css" />
  <link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.13.11/dist/katex.min.css" type="text/css" />
  <link rel="stylesheet" href="_static/katex-math.css" type="text/css" />
  <link rel="stylesheet" href="_static/sphinx-dropdown.css" type="text/css" />
  <link rel="stylesheet" href="_static/panels-bootstrap.min.css" type="text/css" />
  <link rel="stylesheet" href="_static/css/jit.css" type="text/css" />
    <link rel="index" title="Index" href="genindex.html" />
    <link rel="search" title="Search" href="search.html" />
    <link rel="next" title="Automatic differentiation package - torch.autograd" href="autograd.html" />
    <link rel="prev" title="Tensor Views" href="tensor_view.html" />


  <!-- Google Analytics -->
  
    <script async src="https://www.googletagmanager.com/gtag/js?id=UA-117752657-2"></script>
    <script>
      window.dataLayer = window.dataLayer || [];
      function gtag(){dataLayer.push(arguments);}
      gtag('js', new Date());

      gtag('config', 'UA-117752657-2');
    </script>
  
  <!-- End Google Analytics -->
  

  <script src="_static/js/modernizr.min.js"></script>

  <!-- Preload the theme fonts -->

<link rel="preload" href="_static/fonts/FreightSans/freight-sans-book.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="_static/fonts/FreightSans/freight-sans-medium.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="_static/fonts/IBMPlexMono/IBMPlexMono-Medium.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="_static/fonts/FreightSans/freight-sans-bold.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="_static/fonts/FreightSans/freight-sans-medium-italic.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="_static/fonts/IBMPlexMono/IBMPlexMono-SemiBold.woff2" as="font" type="font/woff2" crossorigin="anonymous">

<!-- Preload the katex fonts -->

<link rel="preload" href="https://cdn.jsdelivr.net/npm/katex@0.10.0/dist/fonts/KaTeX_Math-Italic.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="https://cdn.jsdelivr.net/npm/katex@0.10.0/dist/fonts/KaTeX_Main-Regular.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="https://cdn.jsdelivr.net/npm/katex@0.10.0/dist/fonts/KaTeX_Main-Bold.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="https://cdn.jsdelivr.net/npm/katex@0.10.0/dist/fonts/KaTeX_Size1-Regular.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="https://cdn.jsdelivr.net/npm/katex@0.10.0/dist/fonts/KaTeX_Size4-Regular.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="https://cdn.jsdelivr.net/npm/katex@0.10.0/dist/fonts/KaTeX_Size2-Regular.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="https://cdn.jsdelivr.net/npm/katex@0.10.0/dist/fonts/KaTeX_Size3-Regular.woff2" as="font" type="font/woff2" crossorigin="anonymous">
<link rel="preload" href="https://cdn.jsdelivr.net/npm/katex@0.10.0/dist/fonts/KaTeX_Caligraphic-Regular.woff2" as="font" type="font/woff2" crossorigin="anonymous">
  <link rel="stylesheet" href="https://use.fontawesome.com/releases/v5.15.2/css/all.css" integrity="sha384-vSIIfh2YWi9wW0r9iZe7RJPrKwp6bG+s9QZMoITbCckVJqGCCRhc+ccxNcdpHuYu" crossorigin="anonymous">
</head>

<div class="container-fluid header-holder tutorials-header" id="header-holder">
  <div class="container">
    <div class="header-container">
      <a class="header-logo" href="https://pytorch.org/" aria-label="PyTorch"></a>

      <div class="main-menu">
        <ul>
          <li>
            <a href="https://pytorch.org/get-started">Get Started</a>
          </li>

          <li>
            <a href="https://pytorch.org/ecosystem">Ecosystem</a>
          </li>

          <li>
            <a href="https://pytorch.org/mobile">Mobile</a>
          </li>

          <li>
            <a href="https://pytorch.org/blog/">Blog</a>
          </li>

          <li>
            <a href="https://pytorch.org/tutorials">Tutorials</a>
          </li>

          <li class="active docs-active">
            <div id="resourcesDropdownButton" data-toggle="resources-dropdown" class="resources-dropdown">
              <a class="resource-option with-down-orange-arrow">
                Docs
              </a>
              <div class="resources-dropdown-menu">
                <a class="doc-dropdown-option nav-dropdown-item" href="https://pytorch.org/docs/stable/index.html">
                  <span class="dropdown-title">PyTorch</span>
                  <p></p>
                </a>
                <a class="doc-dropdown-option nav-dropdown-item" href="https://pytorch.org/audio/stable/index.html">
                  <span class="dropdown-title">torchaudio</span>
                  <p></p>
                </a>
                <a class="doc-dropdown-option nav-dropdown-item" href="https://pytorch.org/text/stable/index.html">
                  <span class="dropdown-title">torchtext</span>
                  <p></p>
                </a>
                <a class="doc-dropdown-option nav-dropdown-item" href="https://pytorch.org/vision/stable/index.html">
                  <span class="dropdown-title">torchvision</span>
                  <p></p>
                </a>
                <a class="doc-dropdown-option nav-dropdown-item" href="https://pytorch.org/torcharrow">
                  <span class="dropdown-title">torcharrow</span>
                  <p></p>
                </a>
                <a class="doc-dropdown-option nav-dropdown-item" href="https://pytorch.org/data">
                  <span class="dropdown-title">TorchData</span>
                  <p></p>
                </a>
                <a class="doc-dropdown-option nav-dropdown-item" href="https://pytorch.org/torchrec">
                  <span class="dropdown-title">TorchRec</span>
                  <p></p>
                </a>
                <a class="doc-dropdown-option nav-dropdown-item" href="https://pytorch.org/serve/">
                  <span class="dropdown-title">TorchServe</span>
                  <p></p>
                </a>
                <a class="doc-dropdown-option nav-dropdown-item" href="https://pytorch.org/torchx/">
                  <span class="dropdown-title">TorchX</span>
                  <p></p>
                </a>
                <a class="doc-dropdown-option nav-dropdown-item" href="https://pytorch.org/xla">
                  <span class="dropdown-title">PyTorch on XLA Devices</span>
                  <p></p>
                </a>
            </div>
          </li>

          <li>
            <div id="resourcesDropdownButton" data-toggle="resources-dropdown" class="resources-dropdown">
              <a class="resource-option with-down-arrow">
                Resources
              </a>
              <div class="resources-dropdown-menu">
                <a class="nav-dropdown-item" href="https://pytorch.org/features">
                  <span class="dropdown-title">About</span>
                  <p>Learn about PyTorch’s features and capabilities</p>
                </a>
                <a class="nav-dropdown-item" href="https://pytorch.org/foundation">
                  <span class="dropdown-title">PyTorch Foundation</span>
                  <p>Learn about the PyTorch foundation</p>
                </a>
                <a class="nav-dropdown-item" href="https://pytorch.org/#community-module">
                  <span class="dropdown-title">Community</span>
                  <p>Join the PyTorch developer community to contribute, learn, and get your questions answered.</p>
                </a>
                <a class="nav-dropdown-item" href="https://pytorch.org/community-stories">
                  <span class="dropdown-title">Community Stories</span>
                  <p>Learn how our community solves real, everyday machine learning problems with PyTorch.</p>
                </a>
                <a class="nav-dropdown-item" href="https://pytorch.org/resources">
                  <span class="dropdown-title">Developer Resources</span>
                  <p>Find resources and get questions answered</p>
                </a>
                <a class="nav-dropdown-item" href="https://pytorch.org/events">
                  <span class="dropdown-title">Events</span>
                  <p>Find events, webinars, and podcasts</p>
                </a>
                <a class="nav-dropdown-item" href="https://discuss.pytorch.org/" target="_blank">
                  <span class="dropdown-title">Forums</span>
                  <p>A place to discuss PyTorch code, issues, install, research</p>
                </a>
                <a class="nav-dropdown-item" href="https://pytorch.org/hub">
                  <span class="dropdown-title">Models (Beta)</span>
                  <p>Discover, publish, and reuse pre-trained models</p>
                </a>
              </div>
            </div>
          </li>

          <li>
            <a href="https://github.com/pytorch/pytorch">GitHub</a>
          </li>
        </ul>
      </div>

      <a class="main-menu-open-button" href="#" data-behavior="open-mobile-menu"></a>
    </div>
  </div>
</div>

<body class="pytorch-body">

   
    <div class="table-of-contents-link-wrapper">
      <span>Table of Contents</span>
      <a href="#" class="toggle-table-of-contents" data-behavior="toggle-table-of-contents"></a>
    </div>

    <nav data-toggle="wy-nav-shift" class="pytorch-left-menu" id="pytorch-left-menu">
      <div class="pytorch-side-scroll">
        <div class="pytorch-menu pytorch-menu-vertical" data-spy="affix" role="navigation" aria-label="main navigation">
          <div class="pytorch-left-menu-search">
            
    <div class="version">
      <a href='https://pytorch.org/docs/versions.html'>1.13 &#x25BC</a>
    </div>
    

<div role="search">
  <form id="rtd-search-form" class="wy-form" action="search.html" method="get">
    <input type="text" name="q" placeholder="Search Docs" />
    <input type="hidden" name="check_keywords" value="yes" />
    <input type="hidden" name="area" value="default" />
  </form>
</div>

          </div>

          
              <p class="caption" role="heading"><span class="caption-text">Community</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="community/build_ci_governance.html">PyTorch Governance | Build + CI</a></li>
<li class="toctree-l1"><a class="reference internal" href="community/contribution_guide.html">PyTorch Contribution Guide</a></li>
<li class="toctree-l1"><a class="reference internal" href="community/design.html">PyTorch Design Philosophy</a></li>
<li class="toctree-l1"><a class="reference internal" href="community/governance.html">PyTorch Governance | Mechanics</a></li>
<li class="toctree-l1"><a class="reference internal" href="community/persons_of_interest.html">PyTorch Governance | Maintainers</a></li>
</ul>
<p class="caption" role="heading"><span class="caption-text">Developer Notes</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="notes/amp_examples.html">CUDA Automatic Mixed Precision examples</a></li>
<li class="toctree-l1"><a class="reference internal" href="notes/autograd.html">Autograd mechanics</a></li>
<li class="toctree-l1"><a class="reference internal" href="notes/broadcasting.html">Broadcasting semantics</a></li>
<li class="toctree-l1"><a class="reference internal" href="notes/cpu_threading_torchscript_inference.html">CPU threading and TorchScript inference</a></li>
<li class="toctree-l1"><a class="reference internal" href="notes/cuda.html">CUDA semantics</a></li>
<li class="toctree-l1"><a class="reference internal" href="notes/ddp.html">Distributed Data Parallel</a></li>
<li class="toctree-l1"><a class="reference internal" href="notes/extending.html">Extending PyTorch</a></li>
<li class="toctree-l1"><a class="reference internal" href="notes/faq.html">Frequently Asked Questions</a></li>
<li class="toctree-l1"><a class="reference internal" href="notes/gradcheck.html">Gradcheck mechanics</a></li>
<li class="toctree-l1"><a class="reference internal" href="notes/hip.html">HIP (ROCm) semantics</a></li>
<li class="toctree-l1"><a class="reference internal" href="notes/large_scale_deployments.html">Features for large-scale deployments</a></li>
<li class="toctree-l1"><a class="reference internal" href="notes/modules.html">Modules</a></li>
<li class="toctree-l1"><a class="reference internal" href="notes/mps.html">MPS backend</a></li>
<li class="toctree-l1"><a class="reference internal" href="notes/multiprocessing.html">Multiprocessing best practices</a></li>
<li class="toctree-l1"><a class="reference internal" href="notes/numerical_accuracy.html">Numerical accuracy</a></li>
<li class="toctree-l1"><a class="reference internal" href="notes/randomness.html">Reproducibility</a></li>
<li class="toctree-l1"><a class="reference internal" href="notes/serialization.html">Serialization semantics</a></li>
<li class="toctree-l1"><a class="reference internal" href="notes/windows.html">Windows FAQ</a></li>
</ul>
<p class="caption" role="heading"><span class="caption-text">Language Bindings</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="cpp_index.html">C++</a></li>
<li class="toctree-l1"><a class="reference external" href="https://pytorch.org/javadoc/">Javadoc</a></li>
<li class="toctree-l1"><a class="reference internal" href="deploy.html">torch::deploy</a></li>
</ul>
<p class="caption" role="heading"><span class="caption-text">Python API</span></p>
<ul class="current">
<li class="toctree-l1"><a class="reference internal" href="torch.html">torch</a></li>
<li class="toctree-l1"><a class="reference internal" href="nn.html">torch.nn</a></li>
<li class="toctree-l1"><a class="reference internal" href="nn.functional.html">torch.nn.functional</a></li>
<li class="toctree-l1"><a class="reference internal" href="tensors.html">torch.Tensor</a></li>
<li class="toctree-l1"><a class="reference internal" href="tensor_attributes.html">Tensor Attributes</a></li>
<li class="toctree-l1"><a class="reference internal" href="tensor_view.html">Tensor Views</a></li>
<li class="toctree-l1 current"><a class="current reference internal" href="#">torch.amp</a></li>
<li class="toctree-l1"><a class="reference internal" href="autograd.html">torch.autograd</a></li>
<li class="toctree-l1"><a class="reference internal" href="library.html">torch.library</a></li>
<li class="toctree-l1"><a class="reference internal" href="cuda.html">torch.cuda</a></li>
<li class="toctree-l1"><a class="reference internal" href="backends.html">torch.backends</a></li>
<li class="toctree-l1"><a class="reference internal" href="distributed.html">torch.distributed</a></li>
<li class="toctree-l1"><a class="reference internal" href="distributed.algorithms.join.html">torch.distributed.algorithms.join</a></li>
<li class="toctree-l1"><a class="reference internal" href="distributed.elastic.html">torch.distributed.elastic</a></li>
<li class="toctree-l1"><a class="reference internal" href="fsdp.html">torch.distributed.fsdp</a></li>
<li class="toctree-l1"><a class="reference internal" href="distributed.optim.html">torch.distributed.optim</a></li>
<li class="toctree-l1"><a class="reference internal" href="distributions.html">torch.distributions</a></li>
<li class="toctree-l1"><a class="reference internal" href="fft.html">torch.fft</a></li>
<li class="toctree-l1"><a class="reference internal" href="futures.html">torch.futures</a></li>
<li class="toctree-l1"><a class="reference internal" href="fx.html">torch.fx</a></li>
<li class="toctree-l1"><a class="reference internal" href="hub.html">torch.hub</a></li>
<li class="toctree-l1"><a class="reference internal" href="jit.html">torch.jit</a></li>
<li class="toctree-l1"><a class="reference internal" href="linalg.html">torch.linalg</a></li>
<li class="toctree-l1"><a class="reference internal" href="monitor.html">torch.monitor</a></li>
<li class="toctree-l1"><a class="reference internal" href="special.html">torch.special</a></li>
<li class="toctree-l1"><a class="reference internal" href="torch.overrides.html">torch.overrides</a></li>
<li class="toctree-l1"><a class="reference internal" href="package.html">torch.package</a></li>
<li class="toctree-l1"><a class="reference internal" href="profiler.html">torch.profiler</a></li>
<li class="toctree-l1"><a class="reference internal" href="nn.init.html">torch.nn.init</a></li>
<li class="toctree-l1"><a class="reference internal" href="onnx.html">torch.onnx</a></li>
<li class="toctree-l1"><a class="reference internal" href="optim.html">torch.optim</a></li>
<li class="toctree-l1"><a class="reference internal" href="complex_numbers.html">Complex Numbers</a></li>
<li class="toctree-l1"><a class="reference internal" href="ddp_comm_hooks.html">DDP Communication Hooks</a></li>
<li class="toctree-l1"><a class="reference internal" href="pipeline.html">Pipeline Parallelism</a></li>
<li class="toctree-l1"><a class="reference internal" href="quantization.html">Quantization</a></li>
<li class="toctree-l1"><a class="reference internal" href="rpc.html">Distributed RPC Framework</a></li>
<li class="toctree-l1"><a class="reference internal" href="random.html">torch.random</a></li>
<li class="toctree-l1"><a class="reference internal" href="masked.html">torch.masked</a></li>
<li class="toctree-l1"><a class="reference internal" href="nested.html">torch.nested</a></li>
<li class="toctree-l1"><a class="reference internal" href="sparse.html">torch.sparse</a></li>
<li class="toctree-l1"><a class="reference internal" href="storage.html">torch.Storage</a></li>
<li class="toctree-l1"><a class="reference internal" href="testing.html">torch.testing</a></li>
<li class="toctree-l1"><a class="reference internal" href="benchmark_utils.html">torch.utils.benchmark</a></li>
<li class="toctree-l1"><a class="reference internal" href="bottleneck.html">torch.utils.bottleneck</a></li>
<li class="toctree-l1"><a class="reference internal" href="checkpoint.html">torch.utils.checkpoint</a></li>
<li class="toctree-l1"><a class="reference internal" href="cpp_extension.html">torch.utils.cpp_extension</a></li>
<li class="toctree-l1"><a class="reference internal" href="data.html">torch.utils.data</a></li>
<li class="toctree-l1"><a class="reference internal" href="jit_utils.html">torch.utils.jit</a></li>
<li class="toctree-l1"><a class="reference internal" href="dlpack.html">torch.utils.dlpack</a></li>
<li class="toctree-l1"><a class="reference internal" href="mobile_optimizer.html">torch.utils.mobile_optimizer</a></li>
<li class="toctree-l1"><a class="reference internal" href="model_zoo.html">torch.utils.model_zoo</a></li>
<li class="toctree-l1"><a class="reference internal" href="tensorboard.html">torch.utils.tensorboard</a></li>
<li class="toctree-l1"><a class="reference internal" href="type_info.html">Type Info</a></li>
<li class="toctree-l1"><a class="reference internal" href="named_tensor.html">Named Tensors</a></li>
<li class="toctree-l1"><a class="reference internal" href="name_inference.html">Named Tensors operator coverage</a></li>
<li class="toctree-l1"><a class="reference internal" href="config_mod.html">torch.__config__</a></li>
</ul>
<p class="caption" role="heading"><span class="caption-text">Libraries</span></p>
<ul>
<li class="toctree-l1"><a class="reference external" href="https://pytorch.org/audio/stable">torchaudio</a></li>
<li class="toctree-l1"><a class="reference external" href="https://pytorch.org/data">TorchData</a></li>
<li class="toctree-l1"><a class="reference external" href="https://pytorch.org/torchrec">TorchRec</a></li>
<li class="toctree-l1"><a class="reference external" href="https://pytorch.org/serve">TorchServe</a></li>
<li class="toctree-l1"><a class="reference external" href="https://pytorch.org/text/stable">torchtext</a></li>
<li class="toctree-l1"><a class="reference external" href="https://pytorch.org/vision/stable">torchvision</a></li>
<li class="toctree-l1"><a class="reference external" href="http://pytorch.org/xla/">PyTorch on XLA Devices</a></li>
</ul>

            
        </div>
      </div>
    </nav>

    <div class="pytorch-container">
      <div class="pytorch-page-level-bar" id="pytorch-page-level-bar">
        <div class="pytorch-breadcrumbs-wrapper">
          

<div role="navigation" aria-label="breadcrumbs navigation">

  <ul class="pytorch-breadcrumbs">
    
      <li>
        <a href="index.html">
          
            Docs
          
        </a> &gt;
      </li>

        
      <li>Automatic Mixed Precision package - torch.amp</li>
    
    
      <li class="pytorch-breadcrumbs-aside">
        
            
            <a href="_sources/amp.rst.txt" rel="nofollow"><img src="_static/images/view-page-source-icon.svg"></a>
          
        
      </li>
    
  </ul>

  
</div>
        </div>

        <div class="pytorch-shortcuts-wrapper" id="pytorch-shortcuts-wrapper">
          Shortcuts
        </div>
      </div>

      <section data-toggle="wy-nav-shift" id="pytorch-content-wrap" class="pytorch-content-wrap">
        <div class="pytorch-content-left">

        
          <div class="rst-content">
          
            <div role="main" class="main-content" itemscope="itemscope" itemtype="http://schema.org/Article">
             <article itemprop="articleBody" id="pytorch-article" class="pytorch-article">
              
  <section id="automatic-mixed-precision-package-torch-amp">
<h1>Automatic Mixed Precision package - torch.amp<a class="headerlink" href="#automatic-mixed-precision-package-torch-amp" title="Permalink to this heading">¶</a></h1>
<span class="target" id="module-torch.cpu"></span><span class="target" id="module-torch.cpu.amp"></span><span class="target" id="module-torch.cuda.amp"></span><span class="target" id="module-torch.amp"></span><p><a class="reference internal" href="#module-torch.amp" title="torch.amp"><code class="xref py py-class docutils literal notranslate"><span class="pre">torch.amp</span></code></a> provides convenience methods for mixed precision,
where some operations use the <code class="docutils literal notranslate"><span class="pre">torch.float32</span></code> (<code class="docutils literal notranslate"><span class="pre">float</span></code>) datatype and other operations
use lower precision floating point datatype (<code class="docutils literal notranslate"><span class="pre">lower_precision_fp</span></code>): <code class="docutils literal notranslate"><span class="pre">torch.float16</span></code> (<code class="docutils literal notranslate"><span class="pre">half</span></code>) or <code class="docutils literal notranslate"><span class="pre">torch.bfloat16</span></code>. Some ops, like linear layers and convolutions,
are much faster in <code class="docutils literal notranslate"><span class="pre">lower_precision_fp</span></code>. Other ops, like reductions, often require the dynamic
range of <code class="docutils literal notranslate"><span class="pre">float32</span></code>.  Mixed precision tries to match each op to its appropriate datatype.</p>
<p>Ordinarily, “automatic mixed precision training” with datatype of <code class="docutils literal notranslate"><span class="pre">torch.float16</span></code> uses <a class="reference internal" href="#torch.autocast" title="torch.autocast"><code class="xref py py-class docutils literal notranslate"><span class="pre">torch.autocast</span></code></a> and
<a class="reference internal" href="#torch.cuda.amp.GradScaler" title="torch.cuda.amp.GradScaler"><code class="xref py py-class docutils literal notranslate"><span class="pre">torch.cuda.amp.GradScaler</span></code></a> together, as shown in the <a class="reference internal" href="notes/amp_examples.html#amp-examples"><span class="std std-ref">CUDA Automatic Mixed Precision examples</span></a>
and <a class="reference external" href="https://pytorch.org/tutorials/recipes/recipes/amp_recipe.html">CUDA Automatic Mixed Precision recipe</a>.
However, <a class="reference internal" href="#torch.autocast" title="torch.autocast"><code class="xref py py-class docutils literal notranslate"><span class="pre">torch.autocast</span></code></a> and <a class="reference internal" href="#torch.cuda.amp.GradScaler" title="torch.cuda.amp.GradScaler"><code class="xref py py-class docutils literal notranslate"><span class="pre">torch.cuda.amp.GradScaler</span></code></a> are modular, and may be used separately if desired.
As shown in the CPU example section of <a class="reference internal" href="#torch.autocast" title="torch.autocast"><code class="xref py py-class docutils literal notranslate"><span class="pre">torch.autocast</span></code></a>, “automatic mixed precision training/inference” on CPU with
datatype of <code class="docutils literal notranslate"><span class="pre">torch.bfloat16</span></code> only uses <a class="reference internal" href="#torch.autocast" title="torch.autocast"><code class="xref py py-class docutils literal notranslate"><span class="pre">torch.autocast</span></code></a>.</p>
<p>For CUDA and CPU, APIs are also provided separately:</p>
<ul class="simple">
<li><p><code class="docutils literal notranslate"><span class="pre">torch.autocast(&quot;cuda&quot;,</span> <span class="pre">args...)</span></code> is equivalent to <code class="docutils literal notranslate"><span class="pre">torch.cuda.amp.autocast(args...)</span></code>.</p></li>
<li><p><code class="docutils literal notranslate"><span class="pre">torch.autocast(&quot;cpu&quot;,</span> <span class="pre">args...)</span></code> is equivalent to <code class="docutils literal notranslate"><span class="pre">torch.cpu.amp.autocast(args...)</span></code>. For CPU, only lower precision floating point datatype of <code class="docutils literal notranslate"><span class="pre">torch.bfloat16</span></code> is supported for now.</p></li>
</ul>
<nav class="contents local" id="contents">
<ul class="simple">
<li><p><a class="reference internal" href="#autocasting" id="id4">Autocasting</a></p></li>
<li><p><a class="reference internal" href="#gradient-scaling" id="id5">Gradient Scaling</a></p></li>
<li><p><a class="reference internal" href="#autocast-op-reference" id="id6">Autocast Op Reference</a></p>
<ul>
<li><p><a class="reference internal" href="#op-eligibility" id="id7">Op Eligibility</a></p></li>
<li><p><a class="reference internal" href="#cuda-op-specific-behavior" id="id8">CUDA Op-Specific Behavior</a></p>
<ul>
<li><p><a class="reference internal" href="#cuda-ops-that-can-autocast-to-float16" id="id9">CUDA Ops that can autocast to <code class="docutils literal notranslate"><span class="pre">float16</span></code></a></p></li>
<li><p><a class="reference internal" href="#cuda-ops-that-can-autocast-to-float32" id="id10">CUDA Ops that can autocast to <code class="docutils literal notranslate"><span class="pre">float32</span></code></a></p></li>
<li><p><a class="reference internal" href="#cuda-ops-that-promote-to-the-widest-input-type" id="id11">CUDA Ops that promote to the widest input type</a></p></li>
<li><p><a class="reference internal" href="#prefer-binary-cross-entropy-with-logits-over-binary-cross-entropy" id="id12">Prefer <code class="docutils literal notranslate"><span class="pre">binary_cross_entropy_with_logits</span></code> over <code class="docutils literal notranslate"><span class="pre">binary_cross_entropy</span></code></a></p></li>
</ul>
</li>
<li><p><a class="reference internal" href="#cpu-op-specific-behavior" id="id13">CPU Op-Specific Behavior</a></p>
<ul>
<li><p><a class="reference internal" href="#cpu-ops-that-can-autocast-to-bfloat16" id="id14">CPU Ops that can autocast to <code class="docutils literal notranslate"><span class="pre">bfloat16</span></code></a></p></li>
<li><p><a class="reference internal" href="#cpu-ops-that-can-autocast-to-float32" id="id15">CPU Ops that can autocast to <code class="docutils literal notranslate"><span class="pre">float32</span></code></a></p></li>
<li><p><a class="reference internal" href="#cpu-ops-that-promote-to-the-widest-input-type" id="id16">CPU Ops that promote to the widest input type</a></p></li>
</ul>
</li>
</ul>
</li>
</ul>
</nav>
<section id="autocasting">
<span id="id1"></span><h2><a class="toc-backref" href="#id4" role="doc-backlink">Autocasting</a><a class="headerlink" href="#autocasting" title="Permalink to this heading">¶</a></h2>
<dl class="py class">
<dt class="sig sig-object py" id="torch.autocast">
<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">torch.</span></span><span class="sig-name descname"><span class="pre">autocast</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">device_type</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">dtype</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">enabled</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">True</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">cache_enabled</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/amp/autocast_mode.html#autocast"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.autocast" title="Permalink to this definition">¶</a></dt>
<dd><p>Instances of <a class="reference internal" href="#torch.autocast" title="torch.autocast"><code class="xref py py-class docutils literal notranslate"><span class="pre">autocast</span></code></a> serve as context managers or decorators that
allow regions of your script to run in mixed precision.</p>
<p>In these regions, ops run in an op-specific dtype chosen by autocast
to improve performance while maintaining accuracy.
See the <a class="reference internal" href="#autocast-op-reference"><span class="std std-ref">Autocast Op Reference</span></a> for details.</p>
<p>When entering an autocast-enabled region, Tensors may be any type.
You should not call <code class="docutils literal notranslate"><span class="pre">half()</span></code> or <code class="docutils literal notranslate"><span class="pre">bfloat16()</span></code> on your model(s) or inputs when using autocasting.</p>
<p><a class="reference internal" href="#torch.autocast" title="torch.autocast"><code class="xref py py-class docutils literal notranslate"><span class="pre">autocast</span></code></a> should wrap only the forward pass(es) of your network, including the loss
computation(s).  Backward passes under autocast are not recommended.
Backward ops run in the same type that autocast used for corresponding forward ops.</p>
<p>Example for CUDA Devices:</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="c1"># Creates model and optimizer in default precision</span>
<span class="n">model</span> <span class="o">=</span> <span class="n">Net</span><span class="p">()</span><span class="o">.</span><span class="n">cuda</span><span class="p">()</span>
<span class="n">optimizer</span> <span class="o">=</span> <span class="n">optim</span><span class="o">.</span><span class="n">SGD</span><span class="p">(</span><span class="n">model</span><span class="o">.</span><span class="n">parameters</span><span class="p">(),</span> <span class="o">...</span><span class="p">)</span>

<span class="k">for</span> <span class="nb">input</span><span class="p">,</span> <span class="n">target</span> <span class="ow">in</span> <span class="n">data</span><span class="p">:</span>
    <span class="n">optimizer</span><span class="o">.</span><span class="n">zero_grad</span><span class="p">()</span>

    <span class="c1"># Enables autocasting for the forward pass (model + loss)</span>
    <span class="k">with</span> <span class="n">autocast</span><span class="p">():</span>
        <span class="n">output</span> <span class="o">=</span> <span class="n">model</span><span class="p">(</span><span class="nb">input</span><span class="p">)</span>
        <span class="n">loss</span> <span class="o">=</span> <span class="n">loss_fn</span><span class="p">(</span><span class="n">output</span><span class="p">,</span> <span class="n">target</span><span class="p">)</span>

    <span class="c1"># Exits the context manager before backward()</span>
    <span class="n">loss</span><span class="o">.</span><span class="n">backward</span><span class="p">()</span>
    <span class="n">optimizer</span><span class="o">.</span><span class="n">step</span><span class="p">()</span>
</pre></div>
</div>
<p>See the <a class="reference internal" href="notes/amp_examples.html#amp-examples"><span class="std std-ref">CUDA Automatic Mixed Precision examples</span></a> for usage (along with gradient scaling)
in more complex scenarios (e.g., gradient penalty, multiple models/losses, custom autograd functions).</p>
<p><a class="reference internal" href="#torch.autocast" title="torch.autocast"><code class="xref py py-class docutils literal notranslate"><span class="pre">autocast</span></code></a> can also be used as a decorator, e.g., on the <code class="docutils literal notranslate"><span class="pre">forward</span></code> method of your model:</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="k">class</span> <span class="nc">AutocastModel</span><span class="p">(</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">):</span>
    <span class="o">...</span>
    <span class="nd">@autocast</span><span class="p">()</span>
    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="nb">input</span><span class="p">):</span>
        <span class="o">...</span>
</pre></div>
</div>
<p>Floating-point Tensors produced in an autocast-enabled region may be <code class="docutils literal notranslate"><span class="pre">float16</span></code>.
After returning to an autocast-disabled region, using them with floating-point
Tensors of different dtypes may cause type mismatch errors.  If so, cast the Tensor(s)
produced in the autocast region back to <code class="docutils literal notranslate"><span class="pre">float32</span></code> (or other dtype if desired).
If a Tensor from the autocast region is already <code class="docutils literal notranslate"><span class="pre">float32</span></code>, the cast is a no-op,
and incurs no additional overhead.
CUDA Example:</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="c1"># Creates some tensors in default dtype (here assumed to be float32)</span>
<span class="n">a_float32</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">rand</span><span class="p">((</span><span class="mi">8</span><span class="p">,</span> <span class="mi">8</span><span class="p">),</span> <span class="n">device</span><span class="o">=</span><span class="s2">&quot;cuda&quot;</span><span class="p">)</span>
<span class="n">b_float32</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">rand</span><span class="p">((</span><span class="mi">8</span><span class="p">,</span> <span class="mi">8</span><span class="p">),</span> <span class="n">device</span><span class="o">=</span><span class="s2">&quot;cuda&quot;</span><span class="p">)</span>
<span class="n">c_float32</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">rand</span><span class="p">((</span><span class="mi">8</span><span class="p">,</span> <span class="mi">8</span><span class="p">),</span> <span class="n">device</span><span class="o">=</span><span class="s2">&quot;cuda&quot;</span><span class="p">)</span>
<span class="n">d_float32</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">rand</span><span class="p">((</span><span class="mi">8</span><span class="p">,</span> <span class="mi">8</span><span class="p">),</span> <span class="n">device</span><span class="o">=</span><span class="s2">&quot;cuda&quot;</span><span class="p">)</span>

<span class="k">with</span> <span class="n">autocast</span><span class="p">():</span>
    <span class="c1"># torch.mm is on autocast&#39;s list of ops that should run in float16.</span>
    <span class="c1"># Inputs are float32, but the op runs in float16 and produces float16 output.</span>
    <span class="c1"># No manual casts are required.</span>
    <span class="n">e_float16</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">mm</span><span class="p">(</span><span class="n">a_float32</span><span class="p">,</span> <span class="n">b_float32</span><span class="p">)</span>
    <span class="c1"># Also handles mixed input types</span>
    <span class="n">f_float16</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">mm</span><span class="p">(</span><span class="n">d_float32</span><span class="p">,</span> <span class="n">e_float16</span><span class="p">)</span>

<span class="c1"># After exiting autocast, calls f_float16.float() to use with d_float32</span>
<span class="n">g_float32</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">mm</span><span class="p">(</span><span class="n">d_float32</span><span class="p">,</span> <span class="n">f_float16</span><span class="o">.</span><span class="n">float</span><span class="p">())</span>
</pre></div>
</div>
<p>CPU Training Example:</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="c1"># Creates model and optimizer in default precision</span>
<span class="n">model</span> <span class="o">=</span> <span class="n">Net</span><span class="p">()</span>
<span class="n">optimizer</span> <span class="o">=</span> <span class="n">optim</span><span class="o">.</span><span class="n">SGD</span><span class="p">(</span><span class="n">model</span><span class="o">.</span><span class="n">parameters</span><span class="p">(),</span> <span class="o">...</span><span class="p">)</span>

<span class="k">for</span> <span class="n">epoch</span> <span class="ow">in</span> <span class="n">epochs</span><span class="p">:</span>
    <span class="k">for</span> <span class="nb">input</span><span class="p">,</span> <span class="n">target</span> <span class="ow">in</span> <span class="n">data</span><span class="p">:</span>
        <span class="n">optimizer</span><span class="o">.</span><span class="n">zero_grad</span><span class="p">()</span>

        <span class="c1"># Runs the forward pass with autocasting.</span>
        <span class="k">with</span> <span class="n">torch</span><span class="o">.</span><span class="n">autocast</span><span class="p">(</span><span class="n">device_type</span><span class="o">=</span><span class="s2">&quot;cpu&quot;</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="n">torch</span><span class="o">.</span><span class="n">bfloat16</span><span class="p">):</span>
            <span class="n">output</span> <span class="o">=</span> <span class="n">model</span><span class="p">(</span><span class="nb">input</span><span class="p">)</span>
            <span class="n">loss</span> <span class="o">=</span> <span class="n">loss_fn</span><span class="p">(</span><span class="n">output</span><span class="p">,</span> <span class="n">target</span><span class="p">)</span>

        <span class="n">loss</span><span class="o">.</span><span class="n">backward</span><span class="p">()</span>
        <span class="n">optimizer</span><span class="o">.</span><span class="n">step</span><span class="p">()</span>
</pre></div>
</div>
<p>CPU Inference Example:</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="c1"># Creates model in default precision</span>
<span class="n">model</span> <span class="o">=</span> <span class="n">Net</span><span class="p">()</span><span class="o">.</span><span class="n">eval</span><span class="p">()</span>

<span class="k">with</span> <span class="n">torch</span><span class="o">.</span><span class="n">autocast</span><span class="p">(</span><span class="n">device_type</span><span class="o">=</span><span class="s2">&quot;cpu&quot;</span><span class="p">,</span> <span class="n">dtype</span><span class="o">=</span><span class="n">torch</span><span class="o">.</span><span class="n">bfloat16</span><span class="p">):</span>
    <span class="k">for</span> <span class="nb">input</span> <span class="ow">in</span> <span class="n">data</span><span class="p">:</span>
        <span class="c1"># Runs the forward pass with autocasting.</span>
        <span class="n">output</span> <span class="o">=</span> <span class="n">model</span><span class="p">(</span><span class="nb">input</span><span class="p">)</span>
</pre></div>
</div>
<p>CPU Inference Example with Jit Trace:</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="k">class</span> <span class="nc">TestModel</span><span class="p">(</span><span class="n">nn</span><span class="o">.</span><span class="n">Module</span><span class="p">):</span>
    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">input_size</span><span class="p">,</span> <span class="n">num_classes</span><span class="p">):</span>
        <span class="nb">super</span><span class="p">(</span><span class="n">TestModel</span><span class="p">,</span> <span class="bp">self</span><span class="p">)</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>
        <span class="bp">self</span><span class="o">.</span><span class="n">fc1</span> <span class="o">=</span> <span class="n">nn</span><span class="o">.</span><span class="n">Linear</span><span class="p">(</span><span class="n">input_size</span><span class="p">,</span> <span class="n">num_classes</span><span class="p">)</span>
    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">x</span><span class="p">):</span>
        <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">fc1</span><span class="p">(</span><span class="n">x</span><span class="p">)</span>

<span class="n">input_size</span> <span class="o">=</span> <span class="mi">2</span>
<span class="n">num_classes</span> <span class="o">=</span> <span class="mi">2</span>
<span class="n">model</span> <span class="o">=</span> <span class="n">TestModel</span><span class="p">(</span><span class="n">input_size</span><span class="p">,</span> <span class="n">num_classes</span><span class="p">)</span><span class="o">.</span><span class="n">eval</span><span class="p">()</span>

<span class="c1"># For now, we suggest to disable the Jit Autocast Pass,</span>
<span class="c1"># As the issue: https://github.com/pytorch/pytorch/issues/75956</span>
<span class="n">torch</span><span class="o">.</span><span class="n">_C</span><span class="o">.</span><span class="n">_jit_set_autocast_mode</span><span class="p">(</span><span class="kc">False</span><span class="p">)</span>

<span class="k">with</span> <span class="n">torch</span><span class="o">.</span><span class="n">cpu</span><span class="o">.</span><span class="n">amp</span><span class="o">.</span><span class="n">autocast</span><span class="p">(</span><span class="n">cache_enabled</span><span class="o">=</span><span class="kc">False</span><span class="p">):</span>
    <span class="n">model</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">jit</span><span class="o">.</span><span class="n">trace</span><span class="p">(</span><span class="n">model</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">randn</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="n">input_size</span><span class="p">))</span>
<span class="n">model</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">jit</span><span class="o">.</span><span class="n">freeze</span><span class="p">(</span><span class="n">model</span><span class="p">)</span>
<span class="c1"># Models Run</span>
<span class="k">for</span> <span class="n">_</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="mi">3</span><span class="p">):</span>
    <span class="n">model</span><span class="p">(</span><span class="n">torch</span><span class="o">.</span><span class="n">randn</span><span class="p">(</span><span class="mi">1</span><span class="p">,</span> <span class="n">input_size</span><span class="p">))</span>
</pre></div>
</div>
<p>Type mismatch errors <em>in</em> an autocast-enabled region are a bug; if this is what you observe,
please file an issue.</p>
<p><code class="docutils literal notranslate"><span class="pre">autocast(enabled=False)</span></code> subregions can be nested in autocast-enabled regions.
Locally disabling autocast can be useful, for example, if you want to force a subregion
to run in a particular <code class="docutils literal notranslate"><span class="pre">dtype</span></code>.  Disabling autocast gives you explicit control over
the execution type.  In the subregion, inputs from the surrounding region
should be cast to <code class="docutils literal notranslate"><span class="pre">dtype</span></code> before use:</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="c1"># Creates some tensors in default dtype (here assumed to be float32)</span>
<span class="n">a_float32</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">rand</span><span class="p">((</span><span class="mi">8</span><span class="p">,</span> <span class="mi">8</span><span class="p">),</span> <span class="n">device</span><span class="o">=</span><span class="s2">&quot;cuda&quot;</span><span class="p">)</span>
<span class="n">b_float32</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">rand</span><span class="p">((</span><span class="mi">8</span><span class="p">,</span> <span class="mi">8</span><span class="p">),</span> <span class="n">device</span><span class="o">=</span><span class="s2">&quot;cuda&quot;</span><span class="p">)</span>
<span class="n">c_float32</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">rand</span><span class="p">((</span><span class="mi">8</span><span class="p">,</span> <span class="mi">8</span><span class="p">),</span> <span class="n">device</span><span class="o">=</span><span class="s2">&quot;cuda&quot;</span><span class="p">)</span>
<span class="n">d_float32</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">rand</span><span class="p">((</span><span class="mi">8</span><span class="p">,</span> <span class="mi">8</span><span class="p">),</span> <span class="n">device</span><span class="o">=</span><span class="s2">&quot;cuda&quot;</span><span class="p">)</span>

<span class="k">with</span> <span class="n">autocast</span><span class="p">():</span>
    <span class="n">e_float16</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">mm</span><span class="p">(</span><span class="n">a_float32</span><span class="p">,</span> <span class="n">b_float32</span><span class="p">)</span>
    <span class="k">with</span> <span class="n">autocast</span><span class="p">(</span><span class="n">enabled</span><span class="o">=</span><span class="kc">False</span><span class="p">):</span>
        <span class="c1"># Calls e_float16.float() to ensure float32 execution</span>
        <span class="c1"># (necessary because e_float16 was created in an autocasted region)</span>
        <span class="n">f_float32</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">mm</span><span class="p">(</span><span class="n">c_float32</span><span class="p">,</span> <span class="n">e_float16</span><span class="o">.</span><span class="n">float</span><span class="p">())</span>

    <span class="c1"># No manual casts are required when re-entering the autocast-enabled region.</span>
    <span class="c1"># torch.mm again runs in float16 and produces float16 output, regardless of input types.</span>
    <span class="n">g_float16</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">mm</span><span class="p">(</span><span class="n">d_float32</span><span class="p">,</span> <span class="n">f_float32</span><span class="p">)</span>
</pre></div>
</div>
<p>The autocast state is thread-local.  If you want it enabled in a new thread, the context manager or decorator
must be invoked in that thread.  This affects <a class="reference internal" href="generated/torch.nn.DataParallel.html#torch.nn.DataParallel" title="torch.nn.DataParallel"><code class="xref py py-class docutils literal notranslate"><span class="pre">torch.nn.DataParallel</span></code></a> and
<a class="reference internal" href="generated/torch.nn.parallel.DistributedDataParallel.html#torch.nn.parallel.DistributedDataParallel" title="torch.nn.parallel.DistributedDataParallel"><code class="xref py py-class docutils literal notranslate"><span class="pre">torch.nn.parallel.DistributedDataParallel</span></code></a> when used with more than one GPU per process
(see <a class="reference internal" href="notes/amp_examples.html#amp-multigpu"><span class="std std-ref">Working with Multiple GPUs</span></a>).</p>
<dl class="field-list simple">
<dt class="field-odd">Parameters<span class="colon">:</span></dt>
<dd class="field-odd"><ul class="simple">
<li><p><strong>device_type</strong> (<a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#str" title="(in Python v3.10)"><em>str</em></a><em>, </em><em>required</em>) – Whether to use ‘cuda’ or ‘cpu’ device</p></li>
<li><p><strong>enabled</strong> (<a class="reference external" href="https://docs.python.org/3/library/functions.html#bool" title="(in Python v3.10)"><em>bool</em></a><em>, </em><em>optional</em>) – Whether autocasting should be enabled in the region.
Default: <code class="docutils literal notranslate"><span class="pre">True</span></code></p></li>
<li><p><strong>dtype</strong> (<em>torch_dtype</em><em>, </em><em>optional</em>) – Whether to use torch.float16 or torch.bfloat16.</p></li>
<li><p><strong>cache_enabled</strong> (<a class="reference external" href="https://docs.python.org/3/library/functions.html#bool" title="(in Python v3.10)"><em>bool</em></a><em>, </em><em>optional</em>) – Whether the weight cache inside autocast should be enabled.
Default: <code class="docutils literal notranslate"><span class="pre">True</span></code></p></li>
</ul>
</dd>
</dl>
</dd></dl>

<dl class="py class">
<dt class="sig sig-object py" id="torch.cuda.amp.autocast">
<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">torch.cuda.amp.</span></span><span class="sig-name descname"><span class="pre">autocast</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">enabled</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">True</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">dtype</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">torch.float16</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">cache_enabled</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">True</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/cuda/amp/autocast_mode.html#autocast"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.cuda.amp.autocast" title="Permalink to this definition">¶</a></dt>
<dd><p>See <a class="reference internal" href="#torch.autocast" title="torch.autocast"><code class="xref py py-class docutils literal notranslate"><span class="pre">torch.autocast</span></code></a>.
<code class="docutils literal notranslate"><span class="pre">torch.cuda.amp.autocast(args...)</span></code> is equivalent to <code class="docutils literal notranslate"><span class="pre">torch.autocast(&quot;cuda&quot;,</span> <span class="pre">args...)</span></code></p>
<dl class="field-list simple">
</dl>
</dd></dl>

<dl class="py function">
<dt class="sig sig-object py" id="torch.cuda.amp.custom_fwd">
<span class="sig-prename descclassname"><span class="pre">torch.cuda.amp.</span></span><span class="sig-name descname"><span class="pre">custom_fwd</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">fwd</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em>, <em class="sig-param"><span class="o"><span class="pre">*</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">cast_inputs</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/cuda/amp/autocast_mode.html#custom_fwd"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.cuda.amp.custom_fwd" title="Permalink to this definition">¶</a></dt>
<dd><p>Helper decorator for <code class="docutils literal notranslate"><span class="pre">forward</span></code> methods of custom autograd functions (subclasses of
<a class="reference internal" href="autograd.html#torch.autograd.Function" title="torch.autograd.Function"><code class="xref py py-class docutils literal notranslate"><span class="pre">torch.autograd.Function</span></code></a>).  See the <a class="reference internal" href="notes/amp_examples.html#amp-custom-examples"><span class="std std-ref">example page</span></a> for more detail.</p>
<dl class="field-list simple">
<dt class="field-odd">Parameters<span class="colon">:</span></dt>
<dd class="field-odd"><p><strong>cast_inputs</strong> (<a class="reference internal" href="tensor_attributes.html#torch.dtype" title="torch.dtype"><code class="xref py py-class docutils literal notranslate"><span class="pre">torch.dtype</span></code></a> or None, optional, default=None) – If not <code class="docutils literal notranslate"><span class="pre">None</span></code>,
when <code class="docutils literal notranslate"><span class="pre">forward</span></code> runs in an autocast-enabled region, casts incoming
floating-point CUDA Tensors to the target dtype (non-floating-point Tensors are not affected),
then executes <code class="docutils literal notranslate"><span class="pre">forward</span></code> with autocast disabled.
If <code class="docutils literal notranslate"><span class="pre">None</span></code>, <code class="docutils literal notranslate"><span class="pre">forward</span></code>’s internal ops execute with the current autocast state.</p>
</dd>
</dl>
<div class="admonition note">
<p class="admonition-title">Note</p>
<p>If the decorated <code class="docutils literal notranslate"><span class="pre">forward</span></code> is called outside an autocast-enabled region,
<a class="reference internal" href="#torch.cuda.amp.custom_fwd" title="torch.cuda.amp.custom_fwd"><code class="xref py py-func docutils literal notranslate"><span class="pre">custom_fwd</span></code></a> is a no-op and <code class="docutils literal notranslate"><span class="pre">cast_inputs</span></code> has no effect.</p>
</div>
</dd></dl>

<dl class="py function">
<dt class="sig sig-object py" id="torch.cuda.amp.custom_bwd">
<span class="sig-prename descclassname"><span class="pre">torch.cuda.amp.</span></span><span class="sig-name descname"><span class="pre">custom_bwd</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">bwd</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/cuda/amp/autocast_mode.html#custom_bwd"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.cuda.amp.custom_bwd" title="Permalink to this definition">¶</a></dt>
<dd><p>Helper decorator for backward methods of custom autograd functions (subclasses of
<a class="reference internal" href="autograd.html#torch.autograd.Function" title="torch.autograd.Function"><code class="xref py py-class docutils literal notranslate"><span class="pre">torch.autograd.Function</span></code></a>).
Ensures that <code class="docutils literal notranslate"><span class="pre">backward</span></code> executes with the same autocast state as <code class="docutils literal notranslate"><span class="pre">forward</span></code>.
See the <a class="reference internal" href="notes/amp_examples.html#amp-custom-examples"><span class="std std-ref">example page</span></a> for more detail.</p>
</dd></dl>

<dl class="py class">
<dt class="sig sig-object py" id="torch.cpu.amp.autocast">
<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">torch.cpu.amp.</span></span><span class="sig-name descname"><span class="pre">autocast</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">enabled</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">True</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">dtype</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">torch.bfloat16</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">cache_enabled</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">True</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/cpu/amp/autocast_mode.html#autocast"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.cpu.amp.autocast" title="Permalink to this definition">¶</a></dt>
<dd><p>See <a class="reference internal" href="#torch.autocast" title="torch.autocast"><code class="xref py py-class docutils literal notranslate"><span class="pre">torch.autocast</span></code></a>.
<code class="docutils literal notranslate"><span class="pre">torch.cpu.amp.autocast(args...)</span></code> is equivalent to <code class="docutils literal notranslate"><span class="pre">torch.autocast(&quot;cpu&quot;,</span> <span class="pre">args...)</span></code></p>
<dl class="field-list simple">
</dl>
</dd></dl>

</section>
<section id="gradient-scaling">
<span id="id2"></span><h2><a class="toc-backref" href="#id5" role="doc-backlink">Gradient Scaling</a><a class="headerlink" href="#gradient-scaling" title="Permalink to this heading">¶</a></h2>
<p>If the forward pass for a particular op has <code class="docutils literal notranslate"><span class="pre">float16</span></code> inputs, the backward pass for
that op will produce <code class="docutils literal notranslate"><span class="pre">float16</span></code> gradients.
Gradient values with small magnitudes may not be representable in <code class="docutils literal notranslate"><span class="pre">float16</span></code>.
These values will flush to zero (“underflow”), so the update for the corresponding parameters will be lost.</p>
<p>To prevent underflow, “gradient scaling” multiplies the network’s loss(es) by a scale factor and
invokes a backward pass on the scaled loss(es).  Gradients flowing backward through the network are
then scaled by the same factor.  In other words, gradient values have a larger magnitude,
so they don’t flush to zero.</p>
<p>Each parameter’s gradient (<code class="docutils literal notranslate"><span class="pre">.grad</span></code> attribute) should be unscaled before the optimizer
updates the parameters, so the scale factor does not interfere with the learning rate.</p>
<dl class="py class">
<dt class="sig sig-object py" id="torch.cuda.amp.GradScaler">
<em class="property"><span class="pre">class</span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">torch.cuda.amp.</span></span><span class="sig-name descname"><span class="pre">GradScaler</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">init_scale</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">65536.0</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">growth_factor</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">2.0</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">backoff_factor</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">0.5</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">growth_interval</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">2000</span></span></em>, <em class="sig-param"><span class="n"><span class="pre">enabled</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">True</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/cuda/amp/grad_scaler.html#GradScaler"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.cuda.amp.GradScaler" title="Permalink to this definition">¶</a></dt>
<dd><dl class="py method">
<dt class="sig sig-object py" id="torch.cuda.amp.GradScaler.get_backoff_factor">
<span class="sig-name descname"><span class="pre">get_backoff_factor</span></span><span class="sig-paren">(</span><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/cuda/amp/grad_scaler.html#GradScaler.get_backoff_factor"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.cuda.amp.GradScaler.get_backoff_factor" title="Permalink to this definition">¶</a></dt>
<dd><p>Returns a Python float containing the scale backoff factor.</p>
</dd></dl>

<dl class="py method">
<dt class="sig sig-object py" id="torch.cuda.amp.GradScaler.get_growth_factor">
<span class="sig-name descname"><span class="pre">get_growth_factor</span></span><span class="sig-paren">(</span><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/cuda/amp/grad_scaler.html#GradScaler.get_growth_factor"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.cuda.amp.GradScaler.get_growth_factor" title="Permalink to this definition">¶</a></dt>
<dd><p>Returns a Python float containing the scale growth factor.</p>
</dd></dl>

<dl class="py method">
<dt class="sig sig-object py" id="torch.cuda.amp.GradScaler.get_growth_interval">
<span class="sig-name descname"><span class="pre">get_growth_interval</span></span><span class="sig-paren">(</span><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/cuda/amp/grad_scaler.html#GradScaler.get_growth_interval"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.cuda.amp.GradScaler.get_growth_interval" title="Permalink to this definition">¶</a></dt>
<dd><p>Returns a Python int containing the growth interval.</p>
</dd></dl>

<dl class="py method">
<dt class="sig sig-object py" id="torch.cuda.amp.GradScaler.get_scale">
<span class="sig-name descname"><span class="pre">get_scale</span></span><span class="sig-paren">(</span><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/cuda/amp/grad_scaler.html#GradScaler.get_scale"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.cuda.amp.GradScaler.get_scale" title="Permalink to this definition">¶</a></dt>
<dd><p>Returns a Python float containing the current scale, or 1.0 if scaling is disabled.</p>
<div class="admonition warning">
<p class="admonition-title">Warning</p>
<p><a class="reference internal" href="#torch.cuda.amp.GradScaler.get_scale" title="torch.cuda.amp.GradScaler.get_scale"><code class="xref py py-meth docutils literal notranslate"><span class="pre">get_scale()</span></code></a> incurs a CPU-GPU sync.</p>
</div>
</dd></dl>

<dl class="py method">
<dt class="sig sig-object py" id="torch.cuda.amp.GradScaler.is_enabled">
<span class="sig-name descname"><span class="pre">is_enabled</span></span><span class="sig-paren">(</span><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/cuda/amp/grad_scaler.html#GradScaler.is_enabled"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.cuda.amp.GradScaler.is_enabled" title="Permalink to this definition">¶</a></dt>
<dd><p>Returns a bool indicating whether this instance is enabled.</p>
</dd></dl>

<dl class="py method">
<dt class="sig sig-object py" id="torch.cuda.amp.GradScaler.load_state_dict">
<span class="sig-name descname"><span class="pre">load_state_dict</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">state_dict</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/cuda/amp/grad_scaler.html#GradScaler.load_state_dict"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.cuda.amp.GradScaler.load_state_dict" title="Permalink to this definition">¶</a></dt>
<dd><p>Loads the scaler state.  If this instance is disabled, <a class="reference internal" href="#torch.cuda.amp.GradScaler.load_state_dict" title="torch.cuda.amp.GradScaler.load_state_dict"><code class="xref py py-meth docutils literal notranslate"><span class="pre">load_state_dict()</span></code></a> is a no-op.</p>
<dl class="field-list simple">
<dt class="field-odd">Parameters<span class="colon">:</span></dt>
<dd class="field-odd"><p><strong>state_dict</strong> (<a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#dict" title="(in Python v3.10)"><em>dict</em></a>) – scaler state.  Should be an object returned from a call to <a class="reference internal" href="#torch.cuda.amp.GradScaler.state_dict" title="torch.cuda.amp.GradScaler.state_dict"><code class="xref py py-meth docutils literal notranslate"><span class="pre">state_dict()</span></code></a>.</p>
</dd>
</dl>
</dd></dl>

<dl class="py method">
<dt class="sig sig-object py" id="torch.cuda.amp.GradScaler.scale">
<span class="sig-name descname"><span class="pre">scale</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">outputs</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/cuda/amp/grad_scaler.html#GradScaler.scale"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.cuda.amp.GradScaler.scale" title="Permalink to this definition">¶</a></dt>
<dd><p>Multiplies (‘scales’) a tensor or list of tensors by the scale factor.</p>
<p>Returns scaled outputs.  If this instance of <a class="reference internal" href="#torch.cuda.amp.GradScaler" title="torch.cuda.amp.GradScaler"><code class="xref py py-class docutils literal notranslate"><span class="pre">GradScaler</span></code></a> is not enabled, outputs are returned
unmodified.</p>
<dl class="field-list simple">
<dt class="field-odd">Parameters<span class="colon">:</span></dt>
<dd class="field-odd"><p><strong>outputs</strong> (<a class="reference internal" href="tensors.html#torch.Tensor" title="torch.Tensor"><em>Tensor</em></a><em> or </em><em>iterable of Tensors</em>) – Outputs to scale.</p>
</dd>
</dl>
</dd></dl>

<dl class="py method">
<dt class="sig sig-object py" id="torch.cuda.amp.GradScaler.set_backoff_factor">
<span class="sig-name descname"><span class="pre">set_backoff_factor</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">new_factor</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/cuda/amp/grad_scaler.html#GradScaler.set_backoff_factor"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.cuda.amp.GradScaler.set_backoff_factor" title="Permalink to this definition">¶</a></dt>
<dd><dl class="field-list simple">
<dt class="field-odd">Parameters<span class="colon">:</span></dt>
<dd class="field-odd"><p><strong>new_scale</strong> (<a class="reference external" href="https://docs.python.org/3/library/functions.html#float" title="(in Python v3.10)"><em>float</em></a>) – Value to use as the new scale backoff factor.</p>
</dd>
</dl>
</dd></dl>

<dl class="py method">
<dt class="sig sig-object py" id="torch.cuda.amp.GradScaler.set_growth_factor">
<span class="sig-name descname"><span class="pre">set_growth_factor</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">new_factor</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/cuda/amp/grad_scaler.html#GradScaler.set_growth_factor"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.cuda.amp.GradScaler.set_growth_factor" title="Permalink to this definition">¶</a></dt>
<dd><dl class="field-list simple">
<dt class="field-odd">Parameters<span class="colon">:</span></dt>
<dd class="field-odd"><p><strong>new_scale</strong> (<a class="reference external" href="https://docs.python.org/3/library/functions.html#float" title="(in Python v3.10)"><em>float</em></a>) – Value to use as the new scale growth factor.</p>
</dd>
</dl>
</dd></dl>

<dl class="py method">
<dt class="sig sig-object py" id="torch.cuda.amp.GradScaler.set_growth_interval">
<span class="sig-name descname"><span class="pre">set_growth_interval</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">new_interval</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/cuda/amp/grad_scaler.html#GradScaler.set_growth_interval"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.cuda.amp.GradScaler.set_growth_interval" title="Permalink to this definition">¶</a></dt>
<dd><dl class="field-list simple">
<dt class="field-odd">Parameters<span class="colon">:</span></dt>
<dd class="field-odd"><p><strong>new_interval</strong> (<a class="reference external" href="https://docs.python.org/3/library/functions.html#int" title="(in Python v3.10)"><em>int</em></a>) – Value to use as the new growth interval.</p>
</dd>
</dl>
</dd></dl>

<dl class="py method">
<dt class="sig sig-object py" id="torch.cuda.amp.GradScaler.state_dict">
<span class="sig-name descname"><span class="pre">state_dict</span></span><span class="sig-paren">(</span><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/cuda/amp/grad_scaler.html#GradScaler.state_dict"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.cuda.amp.GradScaler.state_dict" title="Permalink to this definition">¶</a></dt>
<dd><p>Returns the state of the scaler as a <a class="reference external" href="https://docs.python.org/3/library/stdtypes.html#dict" title="(in Python v3.10)"><code class="xref py py-class docutils literal notranslate"><span class="pre">dict</span></code></a>.  It contains five entries:</p>
<ul class="simple">
<li><p><code class="docutils literal notranslate"><span class="pre">&quot;scale&quot;</span></code> - a Python float containing the current scale</p></li>
<li><p><code class="docutils literal notranslate"><span class="pre">&quot;growth_factor&quot;</span></code> - a Python float containing the current growth factor</p></li>
<li><p><code class="docutils literal notranslate"><span class="pre">&quot;backoff_factor&quot;</span></code> - a Python float containing the current backoff factor</p></li>
<li><p><code class="docutils literal notranslate"><span class="pre">&quot;growth_interval&quot;</span></code> - a Python int containing the current growth interval</p></li>
<li><p><code class="docutils literal notranslate"><span class="pre">&quot;_growth_tracker&quot;</span></code> - a Python int containing the number of recent consecutive unskipped steps.</p></li>
</ul>
<p>If this instance is not enabled, returns an empty dict.</p>
<div class="admonition note">
<p class="admonition-title">Note</p>
<p>If you wish to checkpoint the scaler’s state after a particular iteration, <a class="reference internal" href="#torch.cuda.amp.GradScaler.state_dict" title="torch.cuda.amp.GradScaler.state_dict"><code class="xref py py-meth docutils literal notranslate"><span class="pre">state_dict()</span></code></a>
should be called after <a class="reference internal" href="#torch.cuda.amp.GradScaler.update" title="torch.cuda.amp.GradScaler.update"><code class="xref py py-meth docutils literal notranslate"><span class="pre">update()</span></code></a>.</p>
</div>
</dd></dl>

<dl class="py method">
<dt class="sig sig-object py" id="torch.cuda.amp.GradScaler.step">
<span class="sig-name descname"><span class="pre">step</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">optimizer</span></span></em>, <em class="sig-param"><span class="o"><span class="pre">*</span></span><span class="n"><span class="pre">args</span></span></em>, <em class="sig-param"><span class="o"><span class="pre">**</span></span><span class="n"><span class="pre">kwargs</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/cuda/amp/grad_scaler.html#GradScaler.step"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.cuda.amp.GradScaler.step" title="Permalink to this definition">¶</a></dt>
<dd><p><a class="reference internal" href="#torch.cuda.amp.GradScaler.step" title="torch.cuda.amp.GradScaler.step"><code class="xref py py-meth docutils literal notranslate"><span class="pre">step()</span></code></a> carries out the following two operations:</p>
<ol class="arabic simple">
<li><p>Internally invokes <code class="docutils literal notranslate"><span class="pre">unscale_(optimizer)</span></code> (unless <a class="reference internal" href="#torch.cuda.amp.GradScaler.unscale_" title="torch.cuda.amp.GradScaler.unscale_"><code class="xref py py-meth docutils literal notranslate"><span class="pre">unscale_()</span></code></a> was explicitly called for <code class="docutils literal notranslate"><span class="pre">optimizer</span></code>
earlier in the iteration).  As part of the <a class="reference internal" href="#torch.cuda.amp.GradScaler.unscale_" title="torch.cuda.amp.GradScaler.unscale_"><code class="xref py py-meth docutils literal notranslate"><span class="pre">unscale_()</span></code></a>, gradients are checked for infs/NaNs.</p></li>
<li><p>If no inf/NaN gradients are found, invokes <code class="docutils literal notranslate"><span class="pre">optimizer.step()</span></code> using the unscaled
gradients.  Otherwise, <code class="docutils literal notranslate"><span class="pre">optimizer.step()</span></code> is skipped to avoid corrupting the params.</p></li>
</ol>
<p><code class="docutils literal notranslate"><span class="pre">*args</span></code> and <code class="docutils literal notranslate"><span class="pre">**kwargs</span></code> are forwarded to <code class="docutils literal notranslate"><span class="pre">optimizer.step()</span></code>.</p>
<p>Returns the return value of <code class="docutils literal notranslate"><span class="pre">optimizer.step(*args,</span> <span class="pre">**kwargs)</span></code>.</p>
<dl class="field-list simple">
<dt class="field-odd">Parameters<span class="colon">:</span></dt>
<dd class="field-odd"><ul class="simple">
<li><p><strong>optimizer</strong> (<a class="reference internal" href="optim.html#torch.optim.Optimizer" title="torch.optim.Optimizer"><em>torch.optim.Optimizer</em></a>) – Optimizer that applies the gradients.</p></li>
<li><p><strong>args</strong> – Any arguments.</p></li>
<li><p><strong>kwargs</strong> – Any keyword arguments.</p></li>
</ul>
</dd>
</dl>
<div class="admonition warning">
<p class="admonition-title">Warning</p>
<p>Closure use is not currently supported.</p>
</div>
</dd></dl>

<dl class="py method">
<dt class="sig sig-object py" id="torch.cuda.amp.GradScaler.unscale_">
<span class="sig-name descname"><span class="pre">unscale_</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">optimizer</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/cuda/amp/grad_scaler.html#GradScaler.unscale_"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.cuda.amp.GradScaler.unscale_" title="Permalink to this definition">¶</a></dt>
<dd><p>Divides (“unscales”) the optimizer’s gradient tensors by the scale factor.</p>
<p><a class="reference internal" href="#torch.cuda.amp.GradScaler.unscale_" title="torch.cuda.amp.GradScaler.unscale_"><code class="xref py py-meth docutils literal notranslate"><span class="pre">unscale_()</span></code></a> is optional, serving cases where you need to
<a class="reference internal" href="notes/amp_examples.html#working-with-unscaled-gradients"><span class="std std-ref">modify or inspect gradients</span></a>
between the backward pass(es) and <a class="reference internal" href="#torch.cuda.amp.GradScaler.step" title="torch.cuda.amp.GradScaler.step"><code class="xref py py-meth docutils literal notranslate"><span class="pre">step()</span></code></a>.
If <a class="reference internal" href="#torch.cuda.amp.GradScaler.unscale_" title="torch.cuda.amp.GradScaler.unscale_"><code class="xref py py-meth docutils literal notranslate"><span class="pre">unscale_()</span></code></a> is not called explicitly,  gradients will be unscaled  automatically during <a class="reference internal" href="#torch.cuda.amp.GradScaler.step" title="torch.cuda.amp.GradScaler.step"><code class="xref py py-meth docutils literal notranslate"><span class="pre">step()</span></code></a>.</p>
<p>Simple example, using <a class="reference internal" href="#torch.cuda.amp.GradScaler.unscale_" title="torch.cuda.amp.GradScaler.unscale_"><code class="xref py py-meth docutils literal notranslate"><span class="pre">unscale_()</span></code></a> to enable clipping of unscaled gradients:</p>
<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="o">...</span>
<span class="n">scaler</span><span class="o">.</span><span class="n">scale</span><span class="p">(</span><span class="n">loss</span><span class="p">)</span><span class="o">.</span><span class="n">backward</span><span class="p">()</span>
<span class="n">scaler</span><span class="o">.</span><span class="n">unscale_</span><span class="p">(</span><span class="n">optimizer</span><span class="p">)</span>
<span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">utils</span><span class="o">.</span><span class="n">clip_grad_norm_</span><span class="p">(</span><span class="n">model</span><span class="o">.</span><span class="n">parameters</span><span class="p">(),</span> <span class="n">max_norm</span><span class="p">)</span>
<span class="n">scaler</span><span class="o">.</span><span class="n">step</span><span class="p">(</span><span class="n">optimizer</span><span class="p">)</span>
<span class="n">scaler</span><span class="o">.</span><span class="n">update</span><span class="p">()</span>
</pre></div>
</div>
<dl class="field-list simple">
<dt class="field-odd">Parameters<span class="colon">:</span></dt>
<dd class="field-odd"><p><strong>optimizer</strong> (<a class="reference internal" href="optim.html#torch.optim.Optimizer" title="torch.optim.Optimizer"><em>torch.optim.Optimizer</em></a>) – Optimizer that owns the gradients to be unscaled.</p>
</dd>
</dl>
<div class="admonition note">
<p class="admonition-title">Note</p>
<p><a class="reference internal" href="#torch.cuda.amp.GradScaler.unscale_" title="torch.cuda.amp.GradScaler.unscale_"><code class="xref py py-meth docutils literal notranslate"><span class="pre">unscale_()</span></code></a> does not incur a CPU-GPU sync.</p>
</div>
<div class="admonition warning">
<p class="admonition-title">Warning</p>
<p><a class="reference internal" href="#torch.cuda.amp.GradScaler.unscale_" title="torch.cuda.amp.GradScaler.unscale_"><code class="xref py py-meth docutils literal notranslate"><span class="pre">unscale_()</span></code></a> should only be called once per optimizer per <a class="reference internal" href="#torch.cuda.amp.GradScaler.step" title="torch.cuda.amp.GradScaler.step"><code class="xref py py-meth docutils literal notranslate"><span class="pre">step()</span></code></a> call,
and only after all gradients for that optimizer’s assigned parameters have been accumulated.
Calling <a class="reference internal" href="#torch.cuda.amp.GradScaler.unscale_" title="torch.cuda.amp.GradScaler.unscale_"><code class="xref py py-meth docutils literal notranslate"><span class="pre">unscale_()</span></code></a> twice for a given optimizer between each <a class="reference internal" href="#torch.cuda.amp.GradScaler.step" title="torch.cuda.amp.GradScaler.step"><code class="xref py py-meth docutils literal notranslate"><span class="pre">step()</span></code></a> triggers a RuntimeError.</p>
</div>
<div class="admonition warning">
<p class="admonition-title">Warning</p>
<p><a class="reference internal" href="#torch.cuda.amp.GradScaler.unscale_" title="torch.cuda.amp.GradScaler.unscale_"><code class="xref py py-meth docutils literal notranslate"><span class="pre">unscale_()</span></code></a> may unscale sparse gradients out of place, replacing the <code class="docutils literal notranslate"><span class="pre">.grad</span></code> attribute.</p>
</div>
</dd></dl>

<dl class="py method">
<dt class="sig sig-object py" id="torch.cuda.amp.GradScaler.update">
<span class="sig-name descname"><span class="pre">update</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">new_scale</span></span><span class="o"><span class="pre">=</span></span><span class="default_value"><span class="pre">None</span></span></em><span class="sig-paren">)</span><a class="reference internal" href="_modules/torch/cuda/amp/grad_scaler.html#GradScaler.update"><span class="viewcode-link"><span class="pre">[source]</span></span></a><a class="headerlink" href="#torch.cuda.amp.GradScaler.update" title="Permalink to this definition">¶</a></dt>
<dd><p>Updates the scale factor.</p>
<p>If any optimizer steps were skipped the scale is multiplied by <code class="docutils literal notranslate"><span class="pre">backoff_factor</span></code>
to reduce it. If <code class="docutils literal notranslate"><span class="pre">growth_interval</span></code> unskipped iterations occurred consecutively,
the scale is multiplied by <code class="docutils literal notranslate"><span class="pre">growth_factor</span></code> to increase it.</p>
<p>Passing <code class="docutils literal notranslate"><span class="pre">new_scale</span></code> sets the new scale value manually. (<code class="docutils literal notranslate"><span class="pre">new_scale</span></code> is not
used directly, it’s used to fill GradScaler’s internal scale tensor. So if
<code class="docutils literal notranslate"><span class="pre">new_scale</span></code> was a tensor, later in-place changes to that tensor will not further
affect the scale GradScaler uses internally.)</p>
<dl class="field-list simple">
<dt class="field-odd">Parameters<span class="colon">:</span></dt>
<dd class="field-odd"><p><strong>new_scale</strong> (float or <code class="xref py py-class docutils literal notranslate"><span class="pre">torch.cuda.FloatTensor</span></code>, optional, default=None) – New scale factor.</p>
</dd>
</dl>
<div class="admonition warning">
<p class="admonition-title">Warning</p>
<p><a class="reference internal" href="#torch.cuda.amp.GradScaler.update" title="torch.cuda.amp.GradScaler.update"><code class="xref py py-meth docutils literal notranslate"><span class="pre">update()</span></code></a> should only be called at the end of the iteration, after <code class="docutils literal notranslate"><span class="pre">scaler.step(optimizer)</span></code> has
been invoked for all optimizers used this iteration.</p>
</div>
</dd></dl>

</dd></dl>

</section>
<section id="autocast-op-reference">
<span id="id3"></span><h2><a class="toc-backref" href="#id6" role="doc-backlink">Autocast Op Reference</a><a class="headerlink" href="#autocast-op-reference" title="Permalink to this heading">¶</a></h2>
<section id="op-eligibility">
<span id="autocast-eligibility"></span><h3><a class="toc-backref" href="#id7" role="doc-backlink">Op Eligibility</a><a class="headerlink" href="#op-eligibility" title="Permalink to this heading">¶</a></h3>
<p>Ops that run in <code class="docutils literal notranslate"><span class="pre">float64</span></code> or non-floating-point dtypes are not eligible, and will
run in these types whether or not autocast is enabled.</p>
<p>Only out-of-place ops and Tensor methods are eligible.
In-place variants and calls that explicitly supply an <code class="docutils literal notranslate"><span class="pre">out=...</span></code> Tensor
are allowed in autocast-enabled regions, but won’t go through autocasting.
For example, in an autocast-enabled region <code class="docutils literal notranslate"><span class="pre">a.addmm(b,</span> <span class="pre">c)</span></code> can autocast,
but <code class="docutils literal notranslate"><span class="pre">a.addmm_(b,</span> <span class="pre">c)</span></code> and <code class="docutils literal notranslate"><span class="pre">a.addmm(b,</span> <span class="pre">c,</span> <span class="pre">out=d)</span></code> cannot.
For best performance and stability, prefer out-of-place ops in autocast-enabled
regions.</p>
<p>Ops called with an explicit <code class="docutils literal notranslate"><span class="pre">dtype=...</span></code> argument are not eligible,
and will produce output that respects the <code class="docutils literal notranslate"><span class="pre">dtype</span></code> argument.</p>
</section>
<section id="cuda-op-specific-behavior">
<span id="autocast-cuda-op-reference"></span><h3><a class="toc-backref" href="#id8" role="doc-backlink">CUDA Op-Specific Behavior</a><a class="headerlink" href="#cuda-op-specific-behavior" title="Permalink to this heading">¶</a></h3>
<p>The following lists describe the behavior of eligible ops in autocast-enabled regions.
These ops always go through autocasting whether they are invoked as part of a <a class="reference internal" href="generated/torch.nn.Module.html#torch.nn.Module" title="torch.nn.Module"><code class="xref py py-class docutils literal notranslate"><span class="pre">torch.nn.Module</span></code></a>,
as a function, or as a <a class="reference internal" href="tensors.html#torch.Tensor" title="torch.Tensor"><code class="xref py py-class docutils literal notranslate"><span class="pre">torch.Tensor</span></code></a> method. If functions are exposed in multiple namespaces,
they go through autocasting regardless of the namespace.</p>
<p>Ops not listed below do not go through autocasting.  They run in the type
defined by their inputs.  However, autocasting may still change the type
in which unlisted ops run if they’re downstream from autocasted ops.</p>
<p>If an op is unlisted, we assume it’s numerically stable in <code class="docutils literal notranslate"><span class="pre">float16</span></code>.
If you believe an unlisted op is numerically unstable in <code class="docutils literal notranslate"><span class="pre">float16</span></code>,
please file an issue.</p>
<section id="cuda-ops-that-can-autocast-to-float16">
<h4><a class="toc-backref" href="#id9" role="doc-backlink">CUDA Ops that can autocast to <code class="docutils literal notranslate"><span class="pre">float16</span></code></a><a class="headerlink" href="#cuda-ops-that-can-autocast-to-float16" title="Permalink to this heading">¶</a></h4>
<p><code class="docutils literal notranslate"><span class="pre">__matmul__</span></code>,
<code class="docutils literal notranslate"><span class="pre">addbmm</span></code>,
<code class="docutils literal notranslate"><span class="pre">addmm</span></code>,
<code class="docutils literal notranslate"><span class="pre">addmv</span></code>,
<code class="docutils literal notranslate"><span class="pre">addr</span></code>,
<code class="docutils literal notranslate"><span class="pre">baddbmm</span></code>,
<code class="docutils literal notranslate"><span class="pre">bmm</span></code>,
<code class="docutils literal notranslate"><span class="pre">chain_matmul</span></code>,
<code class="docutils literal notranslate"><span class="pre">multi_dot</span></code>,
<code class="docutils literal notranslate"><span class="pre">conv1d</span></code>,
<code class="docutils literal notranslate"><span class="pre">conv2d</span></code>,
<code class="docutils literal notranslate"><span class="pre">conv3d</span></code>,
<code class="docutils literal notranslate"><span class="pre">conv_transpose1d</span></code>,
<code class="docutils literal notranslate"><span class="pre">conv_transpose2d</span></code>,
<code class="docutils literal notranslate"><span class="pre">conv_transpose3d</span></code>,
<code class="docutils literal notranslate"><span class="pre">GRUCell</span></code>,
<code class="docutils literal notranslate"><span class="pre">linear</span></code>,
<code class="docutils literal notranslate"><span class="pre">LSTMCell</span></code>,
<code class="docutils literal notranslate"><span class="pre">matmul</span></code>,
<code class="docutils literal notranslate"><span class="pre">mm</span></code>,
<code class="docutils literal notranslate"><span class="pre">mv</span></code>,
<code class="docutils literal notranslate"><span class="pre">prelu</span></code>,
<code class="docutils literal notranslate"><span class="pre">RNNCell</span></code></p>
</section>
<section id="cuda-ops-that-can-autocast-to-float32">
<h4><a class="toc-backref" href="#id10" role="doc-backlink">CUDA Ops that can autocast to <code class="docutils literal notranslate"><span class="pre">float32</span></code></a><a class="headerlink" href="#cuda-ops-that-can-autocast-to-float32" title="Permalink to this heading">¶</a></h4>
<p><code class="docutils literal notranslate"><span class="pre">__pow__</span></code>,
<code class="docutils literal notranslate"><span class="pre">__rdiv__</span></code>,
<code class="docutils literal notranslate"><span class="pre">__rpow__</span></code>,
<code class="docutils literal notranslate"><span class="pre">__rtruediv__</span></code>,
<code class="docutils literal notranslate"><span class="pre">acos</span></code>,
<code class="docutils literal notranslate"><span class="pre">asin</span></code>,
<code class="docutils literal notranslate"><span class="pre">binary_cross_entropy_with_logits</span></code>,
<code class="docutils literal notranslate"><span class="pre">cosh</span></code>,
<code class="docutils literal notranslate"><span class="pre">cosine_embedding_loss</span></code>,
<code class="docutils literal notranslate"><span class="pre">cdist</span></code>,
<code class="docutils literal notranslate"><span class="pre">cosine_similarity</span></code>,
<code class="docutils literal notranslate"><span class="pre">cross_entropy</span></code>,
<code class="docutils literal notranslate"><span class="pre">cumprod</span></code>,
<code class="docutils literal notranslate"><span class="pre">cumsum</span></code>,
<code class="docutils literal notranslate"><span class="pre">dist</span></code>,
<code class="docutils literal notranslate"><span class="pre">erfinv</span></code>,
<code class="docutils literal notranslate"><span class="pre">exp</span></code>,
<code class="docutils literal notranslate"><span class="pre">expm1</span></code>,
<code class="docutils literal notranslate"><span class="pre">group_norm</span></code>,
<code class="docutils literal notranslate"><span class="pre">hinge_embedding_loss</span></code>,
<code class="docutils literal notranslate"><span class="pre">kl_div</span></code>,
<code class="docutils literal notranslate"><span class="pre">l1_loss</span></code>,
<code class="docutils literal notranslate"><span class="pre">layer_norm</span></code>,
<code class="docutils literal notranslate"><span class="pre">log</span></code>,
<code class="docutils literal notranslate"><span class="pre">log_softmax</span></code>,
<code class="docutils literal notranslate"><span class="pre">log10</span></code>,
<code class="docutils literal notranslate"><span class="pre">log1p</span></code>,
<code class="docutils literal notranslate"><span class="pre">log2</span></code>,
<code class="docutils literal notranslate"><span class="pre">margin_ranking_loss</span></code>,
<code class="docutils literal notranslate"><span class="pre">mse_loss</span></code>,
<code class="docutils literal notranslate"><span class="pre">multilabel_margin_loss</span></code>,
<code class="docutils literal notranslate"><span class="pre">multi_margin_loss</span></code>,
<code class="docutils literal notranslate"><span class="pre">nll_loss</span></code>,
<code class="docutils literal notranslate"><span class="pre">norm</span></code>,
<code class="docutils literal notranslate"><span class="pre">normalize</span></code>,
<code class="docutils literal notranslate"><span class="pre">pdist</span></code>,
<code class="docutils literal notranslate"><span class="pre">poisson_nll_loss</span></code>,
<code class="docutils literal notranslate"><span class="pre">pow</span></code>,
<code class="docutils literal notranslate"><span class="pre">prod</span></code>,
<code class="docutils literal notranslate"><span class="pre">reciprocal</span></code>,
<code class="docutils literal notranslate"><span class="pre">rsqrt</span></code>,
<code class="docutils literal notranslate"><span class="pre">sinh</span></code>,
<code class="docutils literal notranslate"><span class="pre">smooth_l1_loss</span></code>,
<code class="docutils literal notranslate"><span class="pre">soft_margin_loss</span></code>,
<code class="docutils literal notranslate"><span class="pre">softmax</span></code>,
<code class="docutils literal notranslate"><span class="pre">softmin</span></code>,
<code class="docutils literal notranslate"><span class="pre">softplus</span></code>,
<code class="docutils literal notranslate"><span class="pre">sum</span></code>,
<code class="docutils literal notranslate"><span class="pre">renorm</span></code>,
<code class="docutils literal notranslate"><span class="pre">tan</span></code>,
<code class="docutils literal notranslate"><span class="pre">triplet_margin_loss</span></code></p>
</section>
<section id="cuda-ops-that-promote-to-the-widest-input-type">
<h4><a class="toc-backref" href="#id11" role="doc-backlink">CUDA Ops that promote to the widest input type</a><a class="headerlink" href="#cuda-ops-that-promote-to-the-widest-input-type" title="Permalink to this heading">¶</a></h4>
<p>These ops don’t require a particular dtype for stability, but take multiple inputs
and require that the inputs’ dtypes match.  If all of the inputs are
<code class="docutils literal notranslate"><span class="pre">float16</span></code>, the op runs in <code class="docutils literal notranslate"><span class="pre">float16</span></code>.  If any of the inputs is <code class="docutils literal notranslate"><span class="pre">float32</span></code>,
autocast casts all inputs to <code class="docutils literal notranslate"><span class="pre">float32</span></code> and runs the op in <code class="docutils literal notranslate"><span class="pre">float32</span></code>.</p>
<p><code class="docutils literal notranslate"><span class="pre">addcdiv</span></code>,
<code class="docutils literal notranslate"><span class="pre">addcmul</span></code>,
<code class="docutils literal notranslate"><span class="pre">atan2</span></code>,
<code class="docutils literal notranslate"><span class="pre">bilinear</span></code>,
<code class="docutils literal notranslate"><span class="pre">cross</span></code>,
<code class="docutils literal notranslate"><span class="pre">dot</span></code>,
<code class="docutils literal notranslate"><span class="pre">grid_sample</span></code>,
<code class="docutils literal notranslate"><span class="pre">index_put</span></code>,
<code class="docutils literal notranslate"><span class="pre">scatter_add</span></code>,
<code class="docutils literal notranslate"><span class="pre">tensordot</span></code></p>
<p>Some ops not listed here (e.g., binary ops like <code class="docutils literal notranslate"><span class="pre">add</span></code>) natively promote
inputs without autocasting’s intervention.  If inputs are a mixture of <code class="docutils literal notranslate"><span class="pre">float16</span></code>
and <code class="docutils literal notranslate"><span class="pre">float32</span></code>, these ops run in <code class="docutils literal notranslate"><span class="pre">float32</span></code> and produce <code class="docutils literal notranslate"><span class="pre">float32</span></code> output,
regardless of whether autocast is enabled.</p>
</section>
<section id="prefer-binary-cross-entropy-with-logits-over-binary-cross-entropy">
<h4><a class="toc-backref" href="#id12" role="doc-backlink">Prefer <code class="docutils literal notranslate"><span class="pre">binary_cross_entropy_with_logits</span></code> over <code class="docutils literal notranslate"><span class="pre">binary_cross_entropy</span></code></a><a class="headerlink" href="#prefer-binary-cross-entropy-with-logits-over-binary-cross-entropy" title="Permalink to this heading">¶</a></h4>
<p>The backward passes of <a class="reference internal" href="generated/torch.nn.functional.binary_cross_entropy.html#torch.nn.functional.binary_cross_entropy" title="torch.nn.functional.binary_cross_entropy"><code class="xref py py-func docutils literal notranslate"><span class="pre">torch.nn.functional.binary_cross_entropy()</span></code></a> (and <a class="reference internal" href="generated/torch.nn.BCELoss.html#torch.nn.BCELoss" title="torch.nn.BCELoss"><code class="xref py py-mod docutils literal notranslate"><span class="pre">torch.nn.BCELoss</span></code></a>, which wraps it)
can produce gradients that aren’t representable in <code class="docutils literal notranslate"><span class="pre">float16</span></code>.  In autocast-enabled regions, the forward input
may be <code class="docutils literal notranslate"><span class="pre">float16</span></code>, which means the backward gradient must be representable in <code class="docutils literal notranslate"><span class="pre">float16</span></code> (autocasting <code class="docutils literal notranslate"><span class="pre">float16</span></code>
forward inputs to <code class="docutils literal notranslate"><span class="pre">float32</span></code> doesn’t help, because that cast must be reversed in backward).
Therefore, <code class="docutils literal notranslate"><span class="pre">binary_cross_entropy</span></code> and <code class="docutils literal notranslate"><span class="pre">BCELoss</span></code> raise an error in autocast-enabled regions.</p>
<p>Many models use a sigmoid layer right before the binary cross entropy layer.
In this case, combine the two layers using <a class="reference internal" href="generated/torch.nn.functional.binary_cross_entropy_with_logits.html#torch.nn.functional.binary_cross_entropy_with_logits" title="torch.nn.functional.binary_cross_entropy_with_logits"><code class="xref py py-func docutils literal notranslate"><span class="pre">torch.nn.functional.binary_cross_entropy_with_logits()</span></code></a>
or <a class="reference internal" href="generated/torch.nn.BCEWithLogitsLoss.html#torch.nn.BCEWithLogitsLoss" title="torch.nn.BCEWithLogitsLoss"><code class="xref py py-mod docutils literal notranslate"><span class="pre">torch.nn.BCEWithLogitsLoss</span></code></a>.  <code class="docutils literal notranslate"><span class="pre">binary_cross_entropy_with_logits</span></code> and <code class="docutils literal notranslate"><span class="pre">BCEWithLogits</span></code>
are safe to autocast.</p>
</section>
</section>
<section id="cpu-op-specific-behavior">
<span id="autocast-cpu-op-reference"></span><h3><a class="toc-backref" href="#id13" role="doc-backlink">CPU Op-Specific Behavior</a><a class="headerlink" href="#cpu-op-specific-behavior" title="Permalink to this heading">¶</a></h3>
<p>The following lists describe the behavior of eligible ops in autocast-enabled regions.
These ops always go through autocasting whether they are invoked as part of a <a class="reference internal" href="generated/torch.nn.Module.html#torch.nn.Module" title="torch.nn.Module"><code class="xref py py-class docutils literal notranslate"><span class="pre">torch.nn.Module</span></code></a>,
as a function, or as a <a class="reference internal" href="tensors.html#torch.Tensor" title="torch.Tensor"><code class="xref py py-class docutils literal notranslate"><span class="pre">torch.Tensor</span></code></a> method. If functions are exposed in multiple namespaces,
they go through autocasting regardless of the namespace.</p>
<p>Ops not listed below do not go through autocasting.  They run in the type
defined by their inputs.  However, autocasting may still change the type
in which unlisted ops run if they’re downstream from autocasted ops.</p>
<p>If an op is unlisted, we assume it’s numerically stable in <code class="docutils literal notranslate"><span class="pre">bfloat16</span></code>.
If you believe an unlisted op is numerically unstable in <code class="docutils literal notranslate"><span class="pre">bfloat16</span></code>,
please file an issue.</p>
<section id="cpu-ops-that-can-autocast-to-bfloat16">
<h4><a class="toc-backref" href="#id14" role="doc-backlink">CPU Ops that can autocast to <code class="docutils literal notranslate"><span class="pre">bfloat16</span></code></a><a class="headerlink" href="#cpu-ops-that-can-autocast-to-bfloat16" title="Permalink to this heading">¶</a></h4>
<p><code class="docutils literal notranslate"><span class="pre">conv1d</span></code>,
<code class="docutils literal notranslate"><span class="pre">conv2d</span></code>,
<code class="docutils literal notranslate"><span class="pre">conv3d</span></code>,
<code class="docutils literal notranslate"><span class="pre">bmm</span></code>,
<code class="docutils literal notranslate"><span class="pre">mm</span></code>,
<code class="docutils literal notranslate"><span class="pre">baddbmm</span></code>,
<code class="docutils literal notranslate"><span class="pre">addmm</span></code>,
<code class="docutils literal notranslate"><span class="pre">addbmm</span></code>,
<code class="docutils literal notranslate"><span class="pre">linear</span></code>,
<code class="docutils literal notranslate"><span class="pre">matmul</span></code>,
<code class="docutils literal notranslate"><span class="pre">_convolution</span></code></p>
</section>
<section id="cpu-ops-that-can-autocast-to-float32">
<h4><a class="toc-backref" href="#id15" role="doc-backlink">CPU Ops that can autocast to <code class="docutils literal notranslate"><span class="pre">float32</span></code></a><a class="headerlink" href="#cpu-ops-that-can-autocast-to-float32" title="Permalink to this heading">¶</a></h4>
<p><code class="docutils literal notranslate"><span class="pre">conv_transpose1d</span></code>,
<code class="docutils literal notranslate"><span class="pre">conv_transpose2d</span></code>,
<code class="docutils literal notranslate"><span class="pre">conv_transpose3d</span></code>,
<code class="docutils literal notranslate"><span class="pre">avg_pool3d</span></code>,
<code class="docutils literal notranslate"><span class="pre">binary_cross_entropy</span></code>,
<code class="docutils literal notranslate"><span class="pre">grid_sampler</span></code>,
<code class="docutils literal notranslate"><span class="pre">grid_sampler_2d</span></code>,
<code class="docutils literal notranslate"><span class="pre">_grid_sampler_2d_cpu_fallback</span></code>,
<code class="docutils literal notranslate"><span class="pre">grid_sampler_3d</span></code>,
<code class="docutils literal notranslate"><span class="pre">polar</span></code>,
<code class="docutils literal notranslate"><span class="pre">prod</span></code>,
<code class="docutils literal notranslate"><span class="pre">quantile</span></code>,
<code class="docutils literal notranslate"><span class="pre">nanquantile</span></code>,
<code class="docutils literal notranslate"><span class="pre">stft</span></code>,
<code class="docutils literal notranslate"><span class="pre">cdist</span></code>,
<code class="docutils literal notranslate"><span class="pre">trace</span></code>,
<code class="docutils literal notranslate"><span class="pre">view_as_complex</span></code>,
<code class="docutils literal notranslate"><span class="pre">cholesky</span></code>,
<code class="docutils literal notranslate"><span class="pre">cholesky_inverse</span></code>,
<code class="docutils literal notranslate"><span class="pre">cholesky_solve</span></code>,
<code class="docutils literal notranslate"><span class="pre">inverse</span></code>,
<code class="docutils literal notranslate"><span class="pre">lu_solve</span></code>,
<code class="docutils literal notranslate"><span class="pre">orgqr</span></code>,
<code class="docutils literal notranslate"><span class="pre">inverse</span></code>,
<code class="docutils literal notranslate"><span class="pre">ormqr</span></code>,
<code class="docutils literal notranslate"><span class="pre">pinverse</span></code>,
<code class="docutils literal notranslate"><span class="pre">max_pool3d</span></code>,
<code class="docutils literal notranslate"><span class="pre">max_unpool2d</span></code>,
<code class="docutils literal notranslate"><span class="pre">max_unpool3d</span></code>,
<code class="docutils literal notranslate"><span class="pre">adaptive_avg_pool3d</span></code>,
<code class="docutils literal notranslate"><span class="pre">reflection_pad1d</span></code>,
<code class="docutils literal notranslate"><span class="pre">reflection_pad2d</span></code>,
<code class="docutils literal notranslate"><span class="pre">replication_pad1d</span></code>,
<code class="docutils literal notranslate"><span class="pre">replication_pad2d</span></code>,
<code class="docutils literal notranslate"><span class="pre">replication_pad3d</span></code>,
<code class="docutils literal notranslate"><span class="pre">mse_loss</span></code>,
<code class="docutils literal notranslate"><span class="pre">ctc_loss</span></code>,
<code class="docutils literal notranslate"><span class="pre">kl_div</span></code>,
<code class="docutils literal notranslate"><span class="pre">multilabel_margin_loss</span></code>,
<code class="docutils literal notranslate"><span class="pre">fft_fft</span></code>,
<code class="docutils literal notranslate"><span class="pre">fft_ifft</span></code>,
<code class="docutils literal notranslate"><span class="pre">fft_fft2</span></code>,
<code class="docutils literal notranslate"><span class="pre">fft_ifft2</span></code>,
<code class="docutils literal notranslate"><span class="pre">fft_fftn</span></code>,
<code class="docutils literal notranslate"><span class="pre">fft_ifftn</span></code>,
<code class="docutils literal notranslate"><span class="pre">fft_rfft</span></code>,
<code class="docutils literal notranslate"><span class="pre">fft_irfft</span></code>,
<code class="docutils literal notranslate"><span class="pre">fft_rfft2</span></code>,
<code class="docutils literal notranslate"><span class="pre">fft_irfft2</span></code>,
<code class="docutils literal notranslate"><span class="pre">fft_rfftn</span></code>,
<code class="docutils literal notranslate"><span class="pre">fft_irfftn</span></code>,
<code class="docutils literal notranslate"><span class="pre">fft_hfft</span></code>,
<code class="docutils literal notranslate"><span class="pre">fft_ihfft</span></code>,
<code class="docutils literal notranslate"><span class="pre">linalg_matrix_norm</span></code>,
<code class="docutils literal notranslate"><span class="pre">linalg_cond</span></code>,
<code class="docutils literal notranslate"><span class="pre">linalg_matrix_rank</span></code>,
<code class="docutils literal notranslate"><span class="pre">linalg_solve</span></code>,
<code class="docutils literal notranslate"><span class="pre">linalg_cholesky</span></code>,
<code class="docutils literal notranslate"><span class="pre">linalg_svdvals</span></code>,
<code class="docutils literal notranslate"><span class="pre">linalg_eigvals</span></code>,
<code class="docutils literal notranslate"><span class="pre">linalg_eigvalsh</span></code>,
<code class="docutils literal notranslate"><span class="pre">linalg_inv</span></code>,
<code class="docutils literal notranslate"><span class="pre">linalg_householder_product</span></code>,
<code class="docutils literal notranslate"><span class="pre">linalg_tensorinv</span></code>,
<code class="docutils literal notranslate"><span class="pre">linalg_tensorsolve</span></code>,
<code class="docutils literal notranslate"><span class="pre">fake_quantize_per_tensor_affine</span></code>,
<code class="docutils literal notranslate"><span class="pre">eig</span></code>,
<code class="docutils literal notranslate"><span class="pre">geqrf</span></code>,
<code class="docutils literal notranslate"><span class="pre">lstsq</span></code>,
<code class="docutils literal notranslate"><span class="pre">_lu_with_info</span></code>,
<code class="docutils literal notranslate"><span class="pre">qr</span></code>,
<code class="docutils literal notranslate"><span class="pre">solve</span></code>,
<code class="docutils literal notranslate"><span class="pre">svd</span></code>,
<code class="docutils literal notranslate"><span class="pre">symeig</span></code>,
<code class="docutils literal notranslate"><span class="pre">triangular_solve</span></code>,
<code class="docutils literal notranslate"><span class="pre">fractional_max_pool2d</span></code>,
<code class="docutils literal notranslate"><span class="pre">fractional_max_pool3d</span></code>,
<code class="docutils literal notranslate"><span class="pre">adaptive_max_pool3d</span></code>,
<code class="docutils literal notranslate"><span class="pre">multilabel_margin_loss_forward</span></code>,
<code class="docutils literal notranslate"><span class="pre">linalg_qr</span></code>,
<code class="docutils literal notranslate"><span class="pre">linalg_cholesky_ex</span></code>,
<code class="docutils literal notranslate"><span class="pre">linalg_svd</span></code>,
<code class="docutils literal notranslate"><span class="pre">linalg_eig</span></code>,
<code class="docutils literal notranslate"><span class="pre">linalg_eigh</span></code>,
<code class="docutils literal notranslate"><span class="pre">linalg_lstsq</span></code>,
<code class="docutils literal notranslate"><span class="pre">linalg_inv_ex</span></code></p>
</section>
<section id="cpu-ops-that-promote-to-the-widest-input-type">
<h4><a class="toc-backref" href="#id16" role="doc-backlink">CPU Ops that promote to the widest input type</a><a class="headerlink" href="#cpu-ops-that-promote-to-the-widest-input-type" title="Permalink to this heading">¶</a></h4>
<p>These ops don’t require a particular dtype for stability, but take multiple inputs
and require that the inputs’ dtypes match.  If all of the inputs are
<code class="docutils literal notranslate"><span class="pre">bfloat16</span></code>, the op runs in <code class="docutils literal notranslate"><span class="pre">bfloat16</span></code>.  If any of the inputs is <code class="docutils literal notranslate"><span class="pre">float32</span></code>,
autocast casts all inputs to <code class="docutils literal notranslate"><span class="pre">float32</span></code> and runs the op in <code class="docutils literal notranslate"><span class="pre">float32</span></code>.</p>
<p><code class="docutils literal notranslate"><span class="pre">cat</span></code>,
<code class="docutils literal notranslate"><span class="pre">stack</span></code>,
<code class="docutils literal notranslate"><span class="pre">index_copy</span></code></p>
<p>Some ops not listed here (e.g., binary ops like <code class="docutils literal notranslate"><span class="pre">add</span></code>) natively promote
inputs without autocasting’s intervention.  If inputs are a mixture of <code class="docutils literal notranslate"><span class="pre">bfloat16</span></code>
and <code class="docutils literal notranslate"><span class="pre">float32</span></code>, these ops run in <code class="docutils literal notranslate"><span class="pre">float32</span></code> and produce <code class="docutils literal notranslate"><span class="pre">float32</span></code> output,
regardless of whether autocast is enabled.</p>
</section>
</section>
</section>
</section>


             </article>
             
            </div>
            <footer>
  
    <div class="rst-footer-buttons" role="navigation" aria-label="footer navigation">
      
        <a href="autograd.html" class="btn btn-neutral float-right" title="Automatic differentiation package - torch.autograd" accesskey="n" rel="next">Next <img src="_static/images/chevron-right-orange.svg" class="next-page"></a>
      
      
        <a href="tensor_view.html" class="btn btn-neutral" title="Tensor Views" accesskey="p" rel="prev"><img src="_static/images/chevron-right-orange.svg" class="previous-page"> Previous</a>
      
    </div>
  

    <hr>

  
  <div role="contentinfo">
    <p>
        &copy; Copyright 2022, PyTorch Contributors.

    </p>
  </div>
    
      <div>
        Built with <a href="http://sphinx-doc.org/">Sphinx</a> using a <a href="https://github.com/rtfd/sphinx_rtd_theme">theme</a> provided by <a href="https://readthedocs.org">Read the Docs</a>.
      </div>
     

</footer>

          </div>
        </div>

        <div class="pytorch-content-right" id="pytorch-content-right">
          <div class="pytorch-right-menu" id="pytorch-right-menu">
            <div class="pytorch-side-scroll" id="pytorch-side-scroll-right">
              <ul>
<li><a class="reference internal" href="#">Automatic Mixed Precision package - torch.amp</a><ul>
<li><a class="reference internal" href="#autocasting">Autocasting</a></li>
<li><a class="reference internal" href="#gradient-scaling">Gradient Scaling</a></li>
<li><a class="reference internal" href="#autocast-op-reference">Autocast Op Reference</a><ul>
<li><a class="reference internal" href="#op-eligibility">Op Eligibility</a></li>
<li><a class="reference internal" href="#cuda-op-specific-behavior">CUDA Op-Specific Behavior</a><ul>
<li><a class="reference internal" href="#cuda-ops-that-can-autocast-to-float16">CUDA Ops that can autocast to <code class="docutils literal notranslate"><span class="pre">float16</span></code></a></li>
<li><a class="reference internal" href="#cuda-ops-that-can-autocast-to-float32">CUDA Ops that can autocast to <code class="docutils literal notranslate"><span class="pre">float32</span></code></a></li>
<li><a class="reference internal" href="#cuda-ops-that-promote-to-the-widest-input-type">CUDA Ops that promote to the widest input type</a></li>
<li><a class="reference internal" href="#prefer-binary-cross-entropy-with-logits-over-binary-cross-entropy">Prefer <code class="docutils literal notranslate"><span class="pre">binary_cross_entropy_with_logits</span></code> over <code class="docutils literal notranslate"><span class="pre">binary_cross_entropy</span></code></a></li>
</ul>
</li>
<li><a class="reference internal" href="#cpu-op-specific-behavior">CPU Op-Specific Behavior</a><ul>
<li><a class="reference internal" href="#cpu-ops-that-can-autocast-to-bfloat16">CPU Ops that can autocast to <code class="docutils literal notranslate"><span class="pre">bfloat16</span></code></a></li>
<li><a class="reference internal" href="#cpu-ops-that-can-autocast-to-float32">CPU Ops that can autocast to <code class="docutils literal notranslate"><span class="pre">float32</span></code></a></li>
<li><a class="reference internal" href="#cpu-ops-that-promote-to-the-widest-input-type">CPU Ops that promote to the widest input type</a></li>
</ul>
</li>
</ul>
</li>
</ul>
</li>
</ul>

            </div>
          </div>
        </div>
      </section>
    </div>

  
       <script type="text/javascript" id="documentation_options" data-url_root="./" src="_static/documentation_options.js"></script>
         <script data-url_root="./" id="documentation_options" src="_static/documentation_options.js"></script>
         <script src="_static/jquery.js"></script>
         <script src="_static/underscore.js"></script>
         <script src="_static/_sphinx_javascript_frameworks_compat.js"></script>
         <script src="_static/doctools.js"></script>
         <script src="_static/clipboard.min.js"></script>
         <script src="_static/copybutton.js"></script>
     

  <script type="text/javascript" src="_static/js/vendor/popper.min.js"></script>
  <script type="text/javascript" src="_static/js/vendor/bootstrap.min.js"></script>
  <script src="https://cdnjs.cloudflare.com/ajax/libs/list.js/1.5.0/list.min.js"></script>
  <script type="text/javascript" src="_static/js/theme.js"></script>

  <script type="text/javascript">
      jQuery(function () {
          SphinxRtdTheme.Navigation.enable(true);
      });
  </script>
 
<script script type="text/javascript">
  var collapsedSections = ['Developer Notes', 'Language Bindings', 'Libraries', 'Community'];
</script>

<img height="1" width="1" style="border-style:none;" alt="" src="https://www.googleadservices.com/pagead/conversion/795629140/?label=txkmCPmdtosBENSssfsC&amp;guid=ON&amp;script=0"/>


  <!-- Begin Footer -->

  <div class="container-fluid docs-tutorials-resources" id="docs-tutorials-resources">
    <div class="container">
      <div class="row">
        <div class="col-md-4 text-center">
          <h2>Docs</h2>
          <p>Access comprehensive developer documentation for PyTorch</p>
          <a class="with-right-arrow" href="https://pytorch.org/docs/stable/index.html">View Docs</a>
        </div>

        <div class="col-md-4 text-center">
          <h2>Tutorials</h2>
          <p>Get in-depth tutorials for beginners and advanced developers</p>
          <a class="with-right-arrow" href="https://pytorch.org/tutorials">View Tutorials</a>
        </div>

        <div class="col-md-4 text-center">
          <h2>Resources</h2>
          <p>Find development resources and get your questions answered</p>
          <a class="with-right-arrow" href="https://pytorch.org/resources">View Resources</a>
        </div>
      </div>
    </div>
  </div>

  <footer class="site-footer">
    <div class="container footer-container">
      <div class="footer-logo-wrapper">
        <a href="https://pytorch.org/" class="footer-logo"></a>
      </div>

      <div class="footer-links-wrapper">
        <div class="footer-links-col">
          <ul>
            <li class="list-title"><a href="https://pytorch.org/">PyTorch</a></li>
            <li><a href="https://pytorch.org/get-started">Get Started</a></li>
            <li><a href="https://pytorch.org/features">Features</a></li>
            <li><a href="https://pytorch.org/ecosystem">Ecosystem</a></li>
            <li><a href="https://pytorch.org/blog/">Blog</a></li>
            <li><a href="https://github.com/pytorch/pytorch/blob/master/CONTRIBUTING.md">Contributing</a></li>
          </ul>
        </div>

        <div class="footer-links-col">
          <ul>
            <li class="list-title"><a href="https://pytorch.org/resources">Resources</a></li>
            <li><a href="https://pytorch.org/tutorials">Tutorials</a></li>
            <li><a href="https://pytorch.org/docs/stable/index.html">Docs</a></li>
            <li><a href="https://discuss.pytorch.org" target="_blank">Discuss</a></li>
            <li><a href="https://github.com/pytorch/pytorch/issues" target="_blank">Github Issues</a></li>
            <li><a href="https://pytorch.org/assets/brand-guidelines/PyTorch-Brand-Guidelines.pdf" target="_blank">Brand Guidelines</a></li>
          </ul>
        </div>

        <div class="footer-links-col">
          <ul>
            <li class="list-title">Stay up to date</li>
            <li><a href="https://www.facebook.com/pytorch" target="_blank">Facebook</a></li>
            <li><a href="https://twitter.com/pytorch" target="_blank">Twitter</a></li>
            <li><a href="https://www.youtube.com/pytorch" target="_blank">YouTube</a></li>
            <li><a href="https://www.linkedin.com/company/pytorch" target="_blank">LinkedIn</a></li>
          </ul>  
          </div>

        <div class="footer-links-col">
          <ul>
            <li class="list-title">PyTorch Podcasts</li>
            <li><a href="https://open.spotify.com/show/6UzHKeiy368jKfQMKKvJY5" target="_blank">Spotify</a></li>
            <li><a href="https://podcasts.apple.com/us/podcast/pytorch-developer-podcast/id1566080008" target="_blank">Apple</a></li>
            <li><a href="https://www.google.com/podcasts?feed=aHR0cHM6Ly9mZWVkcy5zaW1wbGVjYXN0LmNvbS9PQjVGa0lsOA%3D%3D" target="_blank">Google</a></li>
            <li><a href="https://music.amazon.com/podcasts/7a4e6f0e-26c2-49e9-a478-41bd244197d0/PyTorch-Developer-Podcast?" target="_blank">Amazon</a></li>
          </ul>
         </div>
        </div>
        
        <div class="privacy-policy">
          <ul>
            <li class="privacy-policy-links"><a href="https://www.linuxfoundation.org/terms/" target="_blank">Terms</a></li>
            <li class="privacy-policy-links">|</li>
            <li class="privacy-policy-links"><a href="https://www.linuxfoundation.org/privacy-policy/" target="_blank">Privacy</a></li>
          </ul>
        </div>
        <div class="copyright">
        <p>© Copyright The Linux Foundation. The PyTorch Foundation is a project of The Linux Foundation.
          For web site terms of use, trademark policy and other policies applicable to The PyTorch Foundation please see
          <a href="www.linuxfoundation.org/policies/">www.linuxfoundation.org/policies/</a>. The PyTorch Foundation supports the PyTorch open source
          project, which has been established as PyTorch Project a Series of LF Projects, LLC. For policies applicable to the PyTorch Project a Series of LF Projects, LLC,
          please see <a href="www.lfprojects.org/policies/">www.lfprojects.org/policies/</a>.</p>
      </div>
     </div>

  </footer>

  <div class="cookie-banner-wrapper">
  <div class="container">
    <p class="gdpr-notice">To analyze traffic and optimize your experience, we serve cookies on this site. By clicking or navigating, you agree to allow our usage of cookies. As the current maintainers of this site, Facebook’s Cookies Policy applies. Learn more, including about available controls: <a href="https://www.facebook.com/policies/cookies/">Cookies Policy</a>.</p>
    <img class="close-button" src="_static/images/pytorch-x.svg">
  </div>
</div>

  <!-- End Footer -->

  <!-- Begin Mobile Menu -->

  <div class="mobile-main-menu">
    <div class="container-fluid">
      <div class="container">
        <div class="mobile-main-menu-header-container">
          <a class="header-logo" href="https://pytorch.org/" aria-label="PyTorch"></a>
          <a class="main-menu-close-button" href="#" data-behavior="close-mobile-menu"></a>
        </div>
      </div>
    </div>

    <div class="mobile-main-menu-links-container">
      <div class="main-menu">
        <ul>
          <li>
            <a href="https://pytorch.org/get-started">Get Started</a>
          </li>

          <li>
            <a href="https://pytorch.org/ecosystem">Ecosystem</a>
          </li>
            
          <li>
            <a href="https://pytorch.org/mobile">Mobile</a>
          </li>

          <li>
            <a href="https://pytorch.org/blog/">Blog</a>
          </li>

          <li>
            <a href="https://pytorch.org/tutorials">Tutorials</a>
          </li>

          <li class="resources-mobile-menu-title" class="active">
            Docs
          </li>

          <ul class="resources-mobile-menu-items">
            <li>
              <a href="https://pytorch.org/docs/stable/index.html">PyTorch</a>
            </li>

            <li>
              <a href="https://pytorch.org/audio/stable/index.html">torchaudio</a>
            </li>

            <li>
              <a href="https://pytorch.org/text/stable/index.html">torchtext</a>
            </li>

            <li>
              <a href="https://pytorch.org/vision/stable/index.html">torchvision</a>
            </li>

            <li>
              <a href="https://pytorch.org/torcharrow">torcharrow</a>
            </li>

            <li>
              <a href="https://pytorch.org/data">TorchData</a>
            </li>

            <li>
              <a href="https://pytorch.org/torchrec">TorchRec</a>
            </li>

            <li>
              <a href="https://pytorch.org/serve/">TorchServe</a>
            </li>

            <li>
              <a href="https://pytorch.org/torchx/">TorchX</a>
            </li>

            <li>
              <a href="https://pytorch.org/xla">PyTorch on XLA Devices</a>
            </li>
          </ul>

          <li class="resources-mobile-menu-title">
            Resources
          </li>
            
           <ul class="resources-mobile-menu-items">

            <li>
              <a href="https://pytorch.org/features">About</a>
            </li>

            <li>
              <a href="https://pytorch.org/foundation">PyTorch Foundation</a>
            </li>

            <li>
              <a href="https://pytorch.org/#community-module">Community</a>
            </li>

            <li>
              <a href="https://pytorch.org/community-stories">Community Stories</a>
            </li>

            <li>
              <a href="https://pytorch.org/resources">Developer Resources</a>
            </li>

            <li>
              <a href="https://pytorch.org/events">Events</a>
            </li>

            <li>
              <a href="https://discuss.pytorch.org/">Forums</a>
            </li>

            <li>
              <a href="https://pytorch.org/hub">Models (Beta)</a>
            </li>
          </ul>

          <li>
            <a href="https://github.com/pytorch/pytorch">Github</a>
          </li>
        </ul>
      </div>
    </div>
  </div>

  <!-- End Mobile Menu -->

  <script type="text/javascript" src="_static/js/vendor/anchor.min.js"></script>

  <script type="text/javascript">
    $(document).ready(function() {
      mobileMenu.bind();
      mobileTOC.bind();
      pytorchAnchors.bind();
      sideMenus.bind();
      scrollToAnchor.bind();
      highlightNavigation.bind();
      mainMenuDropdown.bind();
      filterTags.bind();

      // Add class to links that have code blocks, since we cannot create links in code blocks
      $("article.pytorch-article a span.pre").each(function(e) {
        $(this).closest("a").addClass("has-code");
      });
    })
  </script>
</body>
</html>