From 6574e2c86a58c87c702128ac53ebefedb1d0152a Mon Sep 17 00:00:00 2001 From: Steven Palma Date: Mon, 3 Mar 2025 14:40:24 +0100 Subject: [PATCH 1/2] [skip ci] fix(examples): Add Tensor type check --- examples/3_train_policy.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/examples/3_train_policy.py b/examples/3_train_policy.py index cf5d4d3ef0..8edc178726 100644 --- a/examples/3_train_policy.py +++ b/examples/3_train_policy.py @@ -85,7 +85,10 @@ 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, non_blocking=True) if isinstance(v, torch.Tensor) else v) + for k, v in batch.items() + } loss, _ = policy.forward(batch) loss.backward() optimizer.step() From f85f55fe62804667b08e7405e57b31bf8faca6e5 Mon Sep 17 00:00:00 2001 From: Steven Palma Date: Mon, 3 Mar 2025 16:59:22 +0100 Subject: [PATCH 2/2] [skip ci] chore(example): remove non_blocking=True to maintain stability with mps --- examples/3_train_policy.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/examples/3_train_policy.py b/examples/3_train_policy.py index 8edc178726..f6eabbfa40 100644 --- a/examples/3_train_policy.py +++ b/examples/3_train_policy.py @@ -85,10 +85,7 @@ def main(): done = False while not done: for batch in dataloader: - batch = { - k: (v.to(device, non_blocking=True) if isinstance(v, torch.Tensor) else v) - 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()