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

fix: mix precision training on tch backend (#209)

This commit is contained in:
Nathaniel Simard 2023-03-07 09:42:51 -05:00 committed by GitHub
parent be96160065
commit cf7847acb5
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
burn-tch/src/ops/tensor.rs
View File

@ -343,7 +343,7 @@ impl<E: TchElement> TensorOps<TchBackend<E>> for TchBackend<E> {
}
fn to_full_precision<const D: usize>(tensor: &TchTensor<E, D>) -> TchTensor<f32, D> {
let tensor = tensor.tensor.to_kind(E::KIND);
let tensor = tensor.tensor.to_kind(tch::Kind::Float);
TchTensor::new(tensor)
}