From bf6f89a5b5a5526766a53ffa337d65ff13dcb5e8 Mon Sep 17 00:00:00 2001
From: Steven Palma <imstevenpmwork@ieee.org>
Date: Mon, 3 Mar 2025 17:01:04 +0100
Subject: [PATCH] fix(examples): Add Tensor type check (#799)

---
 examples/3_train_policy.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/3_train_policy.py b/examples/3_train_policy.py
index cf5d4d3e..f6eabbfa 100644
--- a/examples/3_train_policy.py
+++ b/examples/3_train_policy.py
@@ -85,7 +85,7 @@ def main():
     done = False
     while not done:
         for batch in dataloader:
-            batch = {k: v.to(device, non_blocking=True) for k, v in batch.items()}
+            batch = {k: (v.to(device) if isinstance(v, torch.Tensor) else v) for k, v in batch.items()}
             loss, _ = policy.forward(batch)
             loss.backward()
             optimizer.step()