Pytorchのcuda:0とcpuの分け方(D4RL/rlkit/CQL)

D4RLデータセットを使ったCQL学習を実行する

https://github.com/aviralkumar2907/CQL　を使っています
CQL/d4rl/examples/cql_mujoco_new.py を動かしたいです

必要な環境
CQL/d4rl/enviroment/linux-gpu-env.yml
全てのバージョンが古かったので、自分で今のバージョンをいれました

https://github.com/deepmind/mujoco/releases
2.1.1が必要です
https://github.com/openai/mujoco-py

terminal
1pip install mujoco_py

key:https://roboti.us/license.html
activationkeyがいります

mujocoの初期設定
https://scrapbox.io/programming-notes/MuJoCo%E3%82%A4%E3%83%B3%E3%82%B9%E3%83%88%E3%83%BC%E3%83%AB

home/.mujoco/mujoco-2.1.1/mujoco-2.1.1/hoge

となっているので一つ上げてください

home/.mujoco/mujoco-2.1.1/hoge

発生している問題・エラーメッセージ

Traceback (most recent call last):
  File "cql_mujoco_new.py", line 194, in <module>
    experiment(variant)
  File "cql_mujoco_new.py", line 104, in experiment
    algorithm.train()
  File "/home/kodama/CQL-master/d4rl/rlkit/core/rl_algorithm.py", line 46, in train
    self._train()
  File "/home/kodama/CQL-master/d4rl/rlkit/core/batch_rl_algorithm.py", line 170, in _train
    self.trainer.train(train_data)
  File "/home/kodama/CQL-master/d4rl/rlkit/torch/torch_rl_algorithm.py", line 40, in train
    self.train_from_torch(batch)
  File "/home/kodama/CQL-master/d4rl/rlkit/torch/sac/cql.py", line 237, in train_from_torch
    q1_rand = self._get_tensor_values(obs, random_actions_tensor, network=self.qf1)
  File "/home/kodama/CQL-master/d4rl/rlkit/torch/sac/cql.py", line 135, in _get_tensor_values
    preds = network(obs_temp, actions)
  File "/home/kodama/anaconda3/envs/cql-rlkit/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/kodama/CQL-master/d4rl/rlkit/torch/networks.py", line 87, in forward
    flat_inputs = torch.cat(inputs, dim=1)

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! (when checking argument for argument tensors in method wrapper___cat)

該当のソースコード

terminal
1python ~/CQL-master/d4rl/examples/cql_mujoco_new.py

python
1#CQL-master/d4rl/rlkit/torch/networks.py
2    def forward(self, *inputs, **kwargs):
3        flat_inputs = torch.cat(inputs, dim=1)
4        # flat_inputs = torch.cat(inputs, dim=1).to(device)試したが失敗
5        return super().forward(flat_inputs, **kwargs)

試したこと

最初Pytorchとcudatoolkitがうまく繋がっていないのかと思いましたが、エラーからだと認識しているように見えます。
様々なdeviceの指定などがうまくできていないように思います。

補足情報（FW/ツールのバージョンなど）

terminal
1$conda list
2 Name                    Version                   Build  Channel
3.
4.
5.
6cudatoolkit               11.1.74              h6bb024c_0    nvidia
7.
8d4rl                      1.1                       dev_0    <develop>
9.
10dm-control                0.0.416999250            pypi_0    pypi
11dm-env                    1.5                      pypi_0    pypi
12dm-tree                   0.1.6                    pypi_0    pypi
13.
14gtimer                    1.0.0b5                  pypi_0    pypi
15gym                       0.21.0                   pypi_0    pypi
16.
17.
18.
19matplotlib                3.3.1                         0    anaconda
20matplotlib-base           3.3.1            py38h817c723_0    anaconda
21.
22mujoco-py                 2.1.2.14                 pypi_0    pypi
23ncurses                   6.2                  he6710b0_1    anaconda
24numba                     0.51.2           py38h0573a6f_1    anaconda
25numpy                     1.19.1           py38hbc911f0_0    anaconda
26numpy-base                1.19.1           py38hfa32c7d_0    anaconda
27.
28pip                       20.2.4                   py38_0    anaconda
29.
30pybox2d                   2.3.10           py38h709712a_2    conda-forge
31.
32python                    3.8.5                h7579374_1    anaconda
33python-dateutil           2.8.1                      py_0    anaconda
34python_abi                3.8                      2_cp38    conda-forge
35pytorch                   1.10.1          py3.8_cuda11.1_cudnn8.0.5_0    pytorch
36pytorch-mutex             1.0                        cuda    pytorch
37.
38rlkit                     0.2.1.dev0                dev_0    <develop>
39scipy                     1.5.2            py38h0b6359f_0    anaconda
40.
41.
42.