Convert tensor to dtype before moving to device instead of after

thu-ml · Jan 26, 2025 · 0a79016 · 0a79016
1 parent f94437a
commit 0a79016
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tianshou/data/utils/converter.py b/tianshou/data/utils/converter.py
@@ -45,10 +45,10 @@ def to_torch(
         x.dtype.type,
         np.bool_ | np.number,
     ):  # most often case
-        x = torch.from_numpy(x).to(device)
+        x = torch.from_numpy(x)
         if dtype is not None:
             x = x.type(dtype)
-        return x
+        return x.to(device)
     if isinstance(x, torch.Tensor):  # second often case
         if dtype is not None:
             x = x.type(dtype)