Fix output tensor dtype (#1938)

tracel-ai · Jul 1, 2024 · e753b0c · e753b0c
1 parent f15896d
commit e753b0c
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/crates/burn-fusion/src/ops/float.rs b/crates/burn-fusion/src/ops/float.rs
@@ -218,7 +218,7 @@ impl<B: FusionBackend> FloatTensorOps<Self> for Fusion<B> {
         let stream = tensor.stream;
         let out = tensor
             .client
-            .tensor_uninitialized(tensor.shape.clone(), B::FloatElem::dtype());
+            .tensor_uninitialized(tensor.shape.clone(), B::IntElem::dtype());
 
         let desc = UnaryOperationDescription {
             input: tensor.into_description(),