.. DO NOT EDIT.
.. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY.
.. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE:
.. "recipes/recipes/changing_default_device.py"
.. LINE NUMBERS ARE GIVEN BELOW.

.. only:: html

    .. note::
        :class: sphx-glr-download-link-note

        Click :ref:`here <sphx_glr_download_recipes_recipes_changing_default_device.py>`
        to download the full example code

.. rst-class:: sphx-glr-example-title

.. _sphx_glr_recipes_recipes_changing_default_device.py:


Changing default device
=======================

It is common practice to write PyTorch code in a device-agnostic way,
and then switch between CPU and CUDA depending on what hardware is available.
Typically, to do this you might have used if-statements and ``cuda()`` calls
to do this:

.. note::
   This recipe requires PyTorch 2.0.0 or later.

.. GENERATED FROM PYTHON SOURCE LINES 14-28

.. code-block:: default

    import torch

    USE_CUDA = False

    mod = torch.nn.Linear(20, 30)
    if USE_CUDA:
        mod.cuda()

    device = 'cpu'
    if USE_CUDA:
        device = 'cuda'
    inp = torch.randn(128, 20, device=device)
    print(mod(inp).device)





.. rst-class:: sphx-glr-script-out

 .. code-block:: none

    cpu




.. GENERATED FROM PYTHON SOURCE LINES 29-31

PyTorch now also has a context manager which can take care of the
device transfer automatically. Here is an example:

.. GENERATED FROM PYTHON SOURCE LINES 31-37

.. code-block:: default


    with torch.device('cuda'):
        mod = torch.nn.Linear(20, 30)
        print(mod.weight.device)
        print(mod(torch.randn(128, 20)).device)





.. rst-class:: sphx-glr-script-out

 .. code-block:: none

    cuda:0
    cuda:0




.. GENERATED FROM PYTHON SOURCE LINES 38-39

You can also set it globally like this: 

.. GENERATED FROM PYTHON SOURCE LINES 39-46

.. code-block:: default


    torch.set_default_device('cuda')

    mod = torch.nn.Linear(20, 30)
    print(mod.weight.device)
    print(mod(torch.randn(128, 20)).device)





.. rst-class:: sphx-glr-script-out

 .. code-block:: none

    cuda:0
    cuda:0




.. GENERATED FROM PYTHON SOURCE LINES 47-51

This function imposes a slight performance cost on every Python
call to the torch API (not just factory functions). If this
is causing problems for you, please comment on
`this issue <https://github.com/pytorch/pytorch/issues/92701>`__


.. rst-class:: sphx-glr-timing

   **Total running time of the script:** ( 0 minutes  0.007 seconds)


.. _sphx_glr_download_recipes_recipes_changing_default_device.py:

.. only:: html

  .. container:: sphx-glr-footer sphx-glr-footer-example


    .. container:: sphx-glr-download sphx-glr-download-python

      :download:`Download Python source code: changing_default_device.py <changing_default_device.py>`

    .. container:: sphx-glr-download sphx-glr-download-jupyter

      :download:`Download Jupyter notebook: changing_default_device.ipynb <changing_default_device.ipynb>`


.. only:: html

 .. rst-class:: sphx-glr-signature

    `Gallery generated by Sphinx-Gallery <https://sphinx-gallery.github.io>`_