Fix test_examples

2024-04-16 14:07:16 +01:00
parent 9c2f10bd04
commit 43a614c173
2 changed files with 4 additions and 2 deletions
--- a/examples/3_train_policy.py
+++ b/examples/3_train_policy.py
@@ -53,7 +53,8 @@ step = 0
 done = False
 while not done:
    for batch in dataloader:
-        batch = {k: v.to(device, non_blocking=True) for k, v in batch.items()}
+        for k in batch:
+            batch[k] = batch[k].to(device, non_blocking=True)
        info = policy(batch)
        if step % log_freq == 0:
            num_samples = (step + 1) * cfg.batch_size