fix(examples): Add Tensor type check (#799)

This commit is contained in:
Steven Palma 2025-03-03 17:01:04 +01:00 committed by GitHub
parent 8861546ad8
commit bf6f89a5b5
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 1 additions and 1 deletions

View File

@ -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()