From 3361908fc5fba694cbd949f31e2e5bca1daabe19 Mon Sep 17 00:00:00 2001 From: yintong-lu Date: Tue, 15 Oct 2024 14:58:52 +0000 Subject: [PATCH] torch/ao/quantization/utils.py: Moving eps to targeted device to avoid device mismatch issue (#135204) MOTIVATION We recently verified some quantization tests on devices other than cpu (eg. CUDA and Intel Gaudi devices identified as 'hpu'). We noticed a device mismatch error as eps is a tensor created on cpu but other tensors (min_val_neg, max_val_pos, scale, zero_point) are moved to the targeted _device_. CHANGES Move eps to _device_ of other tensors. Pull Request resolved: https://github.com/pytorch/pytorch/pull/135204 Approved by: https://github.com/jgong5, https://github.com/jerryzh168 --- torch/ao/quantization/utils.py | 1 + 1 file changed, 1 insertion(+) diff --git a/torch/ao/quantization/utils.py b/torch/ao/quantization/utils.py index ee497e38c39..89735523c0b 100644 --- a/torch/ao/quantization/utils.py +++ b/torch/ao/quantization/utils.py @@ -649,6 +649,7 @@ def determine_qparams( device = min_val_neg.device scale = torch.ones(min_val_neg.size(), dtype=torch.double, device=device) zero_point = torch.zeros(min_val_neg.size(), dtype=torch.int64, device=device) + eps = eps.to(device) if qscheme == torch.per_tensor_symmetric or qscheme == torch.per_channel_symmetric: max_val_pos = torch.max(-min_val_neg, max_val_pos)