From 93d137ddd99d3435b28af87aa10c80d73524e916 Mon Sep 17 00:00:00 2001
From: Guillaume Lagrange <lagrange.guillaume.1@gmail.com>
Date: Tue, 24 Sep 2024 15:51:04 -0400
Subject: [PATCH] Add comment

---
 crates/burn-fusion/src/ops/qtensor.rs | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/crates/burn-fusion/src/ops/qtensor.rs b/crates/burn-fusion/src/ops/qtensor.rs
index b9f67cb89..0474c1d47 100644
--- a/crates/burn-fusion/src/ops/qtensor.rs
+++ b/crates/burn-fusion/src/ops/qtensor.rs
@@ -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 {