File tree Expand file tree Collapse file tree 1 file changed +9
-10
lines changed Expand file tree Collapse file tree 1 file changed +9
-10
lines changed Original file line number Diff line number Diff line change 58
58
# Imports and setup
59
59
# -----------------
60
60
#
61
+ # %%
62
+ # .. code-block:: bash
63
+ #
64
+ # %%bash
65
+ # pip3 install torchrl mujoco glfw
61
66
62
67
import torchrl
68
+ import torch
69
+ import tqdm
70
+ from typing import Tuple
63
71
64
72
# sphinx_gallery_start_ignore
65
73
import warnings
66
- from typing import Tuple
67
-
68
74
warnings .filterwarnings ("ignore" )
69
75
# sphinx_gallery_end_ignore
70
76
71
- import torch .cuda
72
- import tqdm
73
-
74
- import torch .multiprocessing
75
-
76
77
###############################################################################
77
78
# 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
- )
79
+ device = torch .device ("cuda:0" if torch .cuda .is_available () else "cpu" )
81
80
collector_device = torch .device ("cpu" ) # Change the device to ``cuda`` to use CUDA
82
81
83
82
###############################################################################
You can’t perform that action at this time.
0 commit comments