dnrops
/
burn
mirror of https://github.com/tracel-ai/burn.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
crates/burn-jit/src/ops/qtensor.rs
View File

@ -103,6 +103,9 @@ where
fn q_to_device(tensor: QuantizedTensor<Self>, device: &Device<Self>) -> QuantizedTensor<Self> {
let mut tensor = tensor;
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
}