Add storage device parameter to replay buffer initialization
- Specify storage device for replay buffer to optimize memory management
This commit is contained in:
committed by
Michel Aractingi
parent
1d4ec50a58
commit
9ea79f8a76
@@ -647,6 +647,7 @@ def train(cfg: DictConfig, out_dir: str | None = None, job_name: str | None = No
|
||||
state_keys=cfg.policy.input_shapes.keys(),
|
||||
action_mask=active_action_dims,
|
||||
action_delta=cfg.env.wrapper.delta_action,
|
||||
storage_device=device,
|
||||
)
|
||||
batch_size: int = batch_size // 2 # We will sample from both replay buffer
|
||||
|
||||
|
||||
Reference in New Issue
Block a user