Failed to run imitation learning example

[jlee52]

Hi,

I encountered an issue when running the imitation learning example with CUDA enabled.
Upon executing launcher.py with the USE_CUDA flag set to True, I received the following errors:

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!
Experiment failed with parameters:
{'env_id': 'HumanoidMuscle.walk', 'n_epochs': 400, 'n_steps_per_epoch': 100000, 'n_steps_per_fit': 1000, 'n_eval_episodes': 10, 'n_epochs_save': 25, 'gamma': 0.99, 'results_dir': './logs/loco_mujoco_evalution_2025-01-11_13-10-37/env_id___HumanoidMuscle.walk', 'use_cuda': True, 'seed': 0}
Traceback (most recent call last):

The only modification I made to launcher.py was enabling the USE_CUDA flag. For context, this example is being run in a virtual environment after a fresh installation of loco-mujoco. I attached full error logs, error_logs.txt.

I would appreciate any advice on how to resolve this issue and successfully run the script with CUDA enabled.

Thank you.