Fix test_examples

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