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

Add comment

This commit is contained in:
Guillaume Lagrange 2024-09-24 15:51:04 -04:00
parent e04c8ad894
commit 93d137ddd9
1 changed files with 4 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
crates/burn-fusion/src/ops/qtensor.rs
View File

@ -69,6 +69,8 @@ impl<B: FusionBackend> QTensorOps<Self> for Fusion<B> {
}
};
// TODO: B::QuantizedElem::dtype()?
// Right now we know that the dtype is u32 for jit backends, but we could possibly
// make this more flexible in the future
let qtensor = client.register_tensor(
handles.pop().unwrap(),
shape.dims,
@ -126,6 +128,8 @@ impl<B: FusionBackend> QTensorOps<Self> for Fusion<B> {
let shape: Vec<usize> = tensor.shape.clone();
// TODO: B::QuantizedElem::dtype()?
// Right now we know that the dtype is u32 for jit backends, but we could possibly
// make this more flexible in the future
let out = tensor.client.tensor_uninitialized(shape, u32::dtype());
let streams = if let Some(offset) = &qparams.offset {