diff options
| -rw-r--r-- | candle-core/src/cpu_backend.rs | 4 | ||||
| -rw-r--r-- | candle-core/tests/pool_tests.rs | 22 |
2 files changed, 23 insertions, 3 deletions
diff --git a/candle-core/src/cpu_backend.rs b/candle-core/src/cpu_backend.rs index fa24c434..6d129680 100644 --- a/candle-core/src/cpu_backend.rs +++ b/candle-core/src/cpu_backend.rs @@ -729,10 +729,10 @@ impl Map1 for UpsampleNearest2D { let scale_h = src_h as f64 / dst_h as f64; let scale_w = src_w as f64 / dst_w as f64; let mut dst = vec![T::zero(); b_sz * c * dst_h * dst_w]; - let src_h_idxs = (0..src_h) + let src_h_idxs = (0..dst_h) .map(|h_idx| usize::min(src_h - 1, (h_idx as f64 * scale_h) as usize)) .collect::<Vec<_>>(); - let src_w_idxs = (0..src_w) + let src_w_idxs = (0..dst_w) .map(|w_idx| usize::min(src_w - 1, (w_idx as f64 * scale_w) as usize)) .collect::<Vec<_>>(); for b_idx in 0..b_sz { diff --git a/candle-core/tests/pool_tests.rs b/candle-core/tests/pool_tests.rs index c8ddef97..73bf7434 100644 --- a/candle-core/tests/pool_tests.rs +++ b/candle-core/tests/pool_tests.rs @@ -1,5 +1,5 @@ mod test_utils; -use candle_core::{Device, Tensor}; +use candle_core::{Device, IndexOp, Tensor}; // https://github.com/huggingface/candle/issues/364 #[test] @@ -59,3 +59,23 @@ fn avg_pool2d_pytorch() -> anyhow::Result<()> { assert_eq!(test_utils::to_vec3_round(pool, 4)?, [[[0.085]], [[0.0078]]]); Ok(()) } + +#[test] +fn upsample_nearest2d() -> anyhow::Result<()> { + let t = Tensor::arange(0f32, 6f32, &Device::Cpu)?.reshape((1, 1, 2, 3))?; + let upsampled = t.upsample_nearest2d(4, 6)?.i(0)?.i(0)?; + assert_eq!( + t.i(0)?.i(0)?.to_vec2::<f32>()?, + [[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]] + ); + assert_eq!( + upsampled.to_vec2::<f32>()?, + [ + [0.0, 0.0, 1.0, 1.0, 2.0, 2.0], + [0.0, 0.0, 1.0, 1.0, 2.0, 2.0], + [3.0, 3.0, 4.0, 4.0, 5.0, 5.0], + [3.0, 3.0, 4.0, 4.0, 5.0, 5.0] + ] + ); + Ok(()) +} |