diff options
| -rw-r--r-- | candle-core/src/dtype.rs | 14 | ||||
| -rw-r--r-- | candle-nn/src/optim.rs | 5 |
2 files changed, 19 insertions, 0 deletions
diff --git a/candle-core/src/dtype.rs b/candle-core/src/dtype.rs index c7a1567f..94ca57d8 100644 --- a/candle-core/src/dtype.rs +++ b/candle-core/src/dtype.rs @@ -67,6 +67,20 @@ impl DType { Self::F64 => 8, } } + + pub fn is_int(&self) -> bool { + match self { + Self::U8 | Self::U32 | Self::I64 => true, + Self::BF16 | Self::F16 | Self::F32 | Self::F64 => false, + } + } + + pub fn is_float(&self) -> bool { + match self { + Self::U8 | Self::U32 | Self::I64 => false, + Self::BF16 | Self::F16 | Self::F32 | Self::F64 => true, + } + } } pub trait WithDType: diff --git a/candle-nn/src/optim.rs b/candle-nn/src/optim.rs index 4294d75e..7704bb48 100644 --- a/candle-nn/src/optim.rs +++ b/candle-nn/src/optim.rs @@ -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(); |