Update dqn_with_rnn_tutorial collector device

vmoens · web-flow · commit ea42806ca8b1 · 2024-02-03T09:00:14.000Z
diff --git a/intermediate_source/dqn_with_rnn_tutorial.py b/intermediate_source/dqn_with_rnn_tutorial.py
@@ -387,7 +387,7 @@
 #   For the sake of efficiency, we're only running a few thousands iterations
 #   here. In a real setting, the total number of frames should be set to 1M.
 #
-collector = SyncDataCollector(env, stoch_policy, frames_per_batch=50, total_frames=200)
+collector = SyncDataCollector(env, stoch_policy, frames_per_batch=50, total_frames=200, device=device)
 rb = TensorDictReplayBuffer(
     storage=LazyMemmapStorage(20_000), batch_size=4, prefetch=10
 )

Original file line number	Diff line number	Diff line change
`@@ -387,7 +387,7 @@`
`387`	`387`	`# For the sake of efficiency, we're only running a few thousands iterations`
`388`	`388`	`# here. In a real setting, the total number of frames should be set to 1M.`
`389`	`389`	`#`
`390`		`-collector = SyncDataCollector(env, stoch_policy, frames_per_batch=50, total_frames=200)`
	`390`	`+collector = SyncDataCollector(env, stoch_policy, frames_per_batch=50, total_frames=200, device=device)`
`391`	`391`	`rb = TensorDictReplayBuffer(`
`392`	`392`	`storage=LazyMemmapStorage(20_000), batch_size=4, prefetch=10`
`393`	`393`	`)`