[Python] Why do I get a segfault when using cuda in sb3 with a custom feature extractor?

I’m encountering a segmentation fault whenever I attempt to train a PPO model on my GPU. I have a custom feature extractor that I’d like to run on CUDA, but as soon as the evaluation callback calls evaluate_policy(), a segmentation fault occurs. Here are the key points:

What I’m Trying to Do:

I’ve configured a custom feature extractor inside my PPO architecture.

I explicitly set the device parameter to "cuda" to leverage GPU acceleration.

Observed Behavior:

Training proceeds until the evaluation callback triggers (via evaluate_policy()).

At that point, Python crashes with a segmentation fault:

Fatal Python error: Segmentation fault
...
File "/path/to/torch/autograd/graph.py", line XXX in _engine_run_backward
...
File "/path/to/stable_baselines3/ppo/ppo.py", line XXX in train
...


What I’ve Tried:

Switching the device back to CPU ("cpu") works normally and does not produce the segfault. Confirmed that my GPU setup and driver versions are correct. Other GPU-based trainings (outside of SB3) run without issue. Reinstalled PyTorch and SB3, but the error persists. I’ve read that PPO in SB3 can run on GPU, though it’s not always recommended. However, due to the complexity of my custom feature extractor, I need GPU acceleration.

model = PPO(
policy="MultiInputPolicy",
learning_rate=4e-5,
n_steps=512,
n_epochs=10,
batch_size=64,
env=train_env,
policy_kwargs=policy_kwargs,
verbose=1,
device="cuda",
)


eval_callback = EvalCallback(
eval_env=val_env,
n_eval_episodes=1,
deterministic=True,
verbose=1,
best_model_save_path=f"logs/{selected_curriculum}/",
log_path=f"logs/{selected_curriculum}/",
)
model.learn(
total_timesteps=3_000_000,
progress_bar=True,
callback=eval_callback,
)

Continue reading...