Fix jit q_to_device

This commit is contained in:
Guillaume Lagrange 2024-09-24 14:58:10 -04:00
parent 561841b439
commit d06533d382
1 changed files with 3 additions and 0 deletions

View File

@ -103,6 +103,9 @@ where
fn q_to_device(tensor: QuantizedTensor<Self>, device: &Device<Self>) -> QuantizedTensor<Self> { fn q_to_device(tensor: QuantizedTensor<Self>, device: &Device<Self>) -> QuantizedTensor<Self> {
let mut tensor = tensor; let mut tensor = tensor;
tensor.qtensor = super::to_device(tensor.qtensor, device); tensor.qtensor = super::to_device(tensor.qtensor, device);
tensor.qparams.scale = super::to_device(tensor.qparams.scale, device);
tensor.qparams.offset = tensor.qparams.offset.map(|x| super::to_device(x, device));
tensor tensor
} }