Skip to content

Commit

Permalink
Apply the cast before the scaling. (#2135)
Browse files Browse the repository at this point in the history
  • Loading branch information
LaurentMazare authored Apr 28, 2024
1 parent 805f3be commit e5c8b88
Showing 1 changed file with 1 addition and 1 deletion.
2 changes: 1 addition & 1 deletion candle-nn/src/ops.rs
Original file line number Diff line number Diff line change
Expand Up @@ -70,7 +70,7 @@ pub fn dropout(xs: &Tensor, drop_p: f32) -> Result<Tensor> {
let rand = Tensor::rand(0f32, 1f32, xs.shape(), xs.device())?;
let scale = 1.0 / (1.0 - drop_p as f64);
let drop_p = Tensor::new(drop_p, xs.device())?.broadcast_as(xs.shape())?;
let mask = (rand.ge(&drop_p)? * scale)?.to_dtype(xs.dtype())?;
let mask = (rand.ge(&drop_p)?.to_dtype(xs.dtype())? * scale)?;
xs * mask
}

Expand Down

0 comments on commit e5c8b88

Please sign in to comment.