fix(examples): Add Tensor type check (#799)
This commit is contained in:
parent
8861546ad8
commit
bf6f89a5b5
|
@ -85,7 +85,7 @@ def main():
|
||||||
done = False
|
done = False
|
||||||
while not done:
|
while not done:
|
||||||
for batch in dataloader:
|
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, _ = policy.forward(batch)
|
||||||
loss.backward()
|
loss.backward()
|
||||||
optimizer.step()
|
optimizer.step()
|
||||||
|
|
Loading…
Reference in New Issue