1 files changed, 4 insertions, 1 deletions
diff --git a/tests/datasets/ShapeDatasets.h b/tests/datasets/ShapeDatasets.h
index 57bb9d01b0..4d75a16e47 100644
--- a/tests/datasets/ShapeDatasets.h
+++ b/tests/datasets/ShapeDatasets.h
@@ -45,7 +45,10 @@ public:
     Small1DShapes()
         : ShapeDataset("Shape",
     {
-        TensorShape{ 256U }
+        TensorShape{ 128U },
+                     TensorShape{ 256U },
+                     TensorShape{ 512U },
+                     TensorShape{ 1024U }
     })
     {
     }