huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-16 02:38:10 +00:00
Code Issues Packages Projects Releases Wiki Activity

Only optimize float tensors. (#1069)

Browse Source
This commit is contained in:
Laurent Mazare
2023-10-10 10:05:41 +02:00
committed by GitHub
parent bc3351bce4
commit 9fea56d28e
2 changed files with 19 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
candle-nn/src/optim.rs
View File

@ -41,6 +41,10 @@ impl Optimizer for SGD {
type Config = f64;
fn new(vars: Vec<Var>, learning_rate: f64) -> Result<Self> {
let vars = vars
.into_iter()
.filter(|var| var.dtype().is_float())
.collect();
Ok(Self {
vars,
learning_rate,
@ -116,6 +120,7 @@ impl Optimizer for AdamW {
fn new(vars: Vec<Var>, params: ParamsAdamW) -> Result<Self> {
let vars = vars
.into_iter()
.filter(|var| var.dtype().is_float())
.map(|var| {
let dtype = var.dtype();
let shape = var.shape();