update unit test

jimexist · Feb 19, 2024 · 15bd436 · 15bd436
1 parent 604719a
commit 15bd436
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/swin_transformer.rs b/src/swin_transformer.rs
@@ -956,7 +956,7 @@ mod test {
         let device = Device::Cpu;
         let vb = VarBuilderArgs::zeros(DType::F32, &device);
         let module = SwinEncoder::new(&Default::default(), vb)?;
-        let x = Tensor::zeros(&[1, 3, 224, 224], DType::F32, &device)?;
+        let x = Tensor::zeros(&[1, 56, 56, 96], DType::F32, &device)?;
         let result = module.forward(&x)?;
         assert_eq!(result.dims(), &[1, 49, 768]);
         Ok(())