during normalize, construct mean and std tensors on the same device as the input tensor (#787)

eb7e273d · Irvin Ho · Francisco Massa · 9d9f48a3 · eb7e273d
Commit eb7e273d authored Mar 09, 2019 by Irvin Ho Committed by Francisco Massa Mar 09, 2019
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 2 deletions

torchvision/transforms/functional.py torchvision/transforms/functional.py +2 -2

No files found.
--- a/torchvision/transforms/functional.py
+++ b/torchvision/transforms/functional.py
@@ -203,8 +203,8 @@ def normalize(tensor, mean, std, inplace=False):
    if not inplace:
        tensor = tensor.clone()

-    mean = torch.tensor(mean, dtype=torch.float32)
-    std = torch.tensor(std, dtype=torch.float32)
+    mean = torch.tensor(mean, dtype=torch.float32, device=tensor.device)
+    std = torch.tensor(std, dtype=torch.float32, device=tensor.device)
    tensor.sub_(mean[:, None, None]).div_(std[:, None, None])
    return tensor