[BugFix] sending data to device for optim steps (#655)

albertbou92 · web-flow · commit e803ffa6f0c7 · 2022-11-11T10:07:03.000Z
diff --git a/torchrl/trainers/helpers/trainers.py b/torchrl/trainers/helpers/trainers.py
@@ -222,6 +222,7 @@ def make_trainer(
             "process_optim_batch",
             BatchSubSampler(batch_size=cfg.batch_size, sub_traj_len=cfg.sub_traj_len),
         )
+        trainer.register_op("process_optim_batch", lambda batch: batch.to(device))
 
     if optim_scheduler is not None:
         trainer.register_op("post_optim", optim_scheduler.step)

Original file line number	Diff line number	Diff line change
`@@ -222,6 +222,7 @@ def make_trainer(`
`222`	`222`	`"process_optim_batch",`
`223`	`223`	`BatchSubSampler(batch_size=cfg.batch_size, sub_traj_len=cfg.sub_traj_len),`
`224`	`224`	`)`
	`225`	`+ trainer.register_op("process_optim_batch", lambda batch: batch.to(device))`
`225`	`226`
`226`	`227`	`if optim_scheduler is not None:`
`227`	`228`	`trainer.register_op("post_optim", optim_scheduler.step)`