diff options
Diffstat (limited to 'tests/datasets/ScatterDataset.h')
-rw-r--r-- | tests/datasets/ScatterDataset.h | 4 |
1 files changed, 2 insertions, 2 deletions
diff --git a/tests/datasets/ScatterDataset.h b/tests/datasets/ScatterDataset.h index d204d17855..f7547ecc94 100644 --- a/tests/datasets/ScatterDataset.h +++ b/tests/datasets/ScatterDataset.h @@ -118,8 +118,8 @@ class Small1DScatterDataset final : public ScatterDataset public: Small1DScatterDataset() { - add_config(TensorShape(6U), TensorShape(6U), TensorShape(6U), TensorShape(6U)); - add_config(TensorShape(10U), TensorShape(2U), TensorShape(2U), TensorShape(10U)); + add_config(TensorShape(6U), TensorShape(6U), TensorShape(1U, 6U), TensorShape(6U)); + add_config(TensorShape(10U), TensorShape(2U), TensorShape(1U, 2U), TensorShape(10U)); } }; } // namespace datasets |