File tree Expand file tree Collapse file tree 1 file changed +8
-10
lines changed
Expand file tree Collapse file tree 1 file changed +8
-10
lines changed Original file line number Diff line number Diff line change 5858# Imports and setup
5959# -----------------
6060#
61+ # .. code-block:: bash
62+ #
63+ # %%bash
64+ # pip3 install torchrl mujoco glfw
6165
6266import torchrl
67+ import torch
68+ import tqdm
69+ from typing import Tuple
6370
6471# sphinx_gallery_start_ignore
6572import warnings
66- from typing import Tuple
67-
6873warnings .filterwarnings ("ignore" )
6974# sphinx_gallery_end_ignore
7075
71- import torch .cuda
72- import tqdm
73-
74- import torch .multiprocessing
75-
7676###############################################################################
7777# We will execute the policy on CUDA if available
78- device = (
79- torch .device ("cpu" ) if torch .cuda .device_count () == 0 else torch .device ("cuda:0" )
80- )
78+ device = torch .device ("cuda:0" if torch .cuda .is_available () else "cpu" )
8179collector_device = torch .device ("cpu" ) # Change the device to ``cuda`` to use CUDA
8280
8381###############################################################################
You can’t perform that action at this time.
0 commit comments