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