did not correctly impact JAX, as it was done after the import. Therefore JAX allocated ~80% of memory of the GPU. When I moved this line of code to the beginning the usage of memory dropped to ~15% (Tested on Nvidia Titan V, HalfCheetah-v4 environment).
Also the import:
from stable_baselines3.common.callbacks import EvalCallback,
Is not used as EvalCallback is imported from sbx.sac.actor_critic_evaluation_callback. It is safer to delete this unnecessary import as it can clash with the other and also some programs e.g. isort when sorting the imports change the order and EvalCallback is then imported from stable_baselines3 instead of sbx.sac.actor_critic_evaluation_callback (that happened in my case when I sorted the imports.)
Hello, Setting the flag:
did not correctly impact JAX, as it was done after the import. Therefore JAX allocated ~80% of memory of the GPU. When I moved this line of code to the beginning the usage of memory dropped to ~15% (Tested on Nvidia Titan V, HalfCheetah-v4 environment).
Also the import:
Is not used as EvalCallback is imported from sbx.sac.actor_critic_evaluation_callback. It is safer to delete this unnecessary import as it can clash with the other and also some programs e.g.
isortwhen sorting the imports change the order and EvalCallback is then imported from stable_baselines3 instead of sbx.sac.actor_critic_evaluation_callback (that happened in my case when I sorted the imports.)