diff --git a/kernels/prim_ops/register_prim_ops.cpp b/kernels/prim_ops/register_prim_ops.cpp index 8607c36204d..4254188d6ee 100644 --- a/kernels/prim_ops/register_prim_ops.cpp +++ b/kernels/prim_ops/register_prim_ops.cpp @@ -124,6 +124,12 @@ static Kernel prim_ops[] = { EValue& out = *stack[1]; executorch::aten::Tensor self_tensor = self.to(); + ET_KERNEL_CHECK_MSG( + context, + self_tensor.numel() >= 1, + InvalidArgument, + /* void */, + "Expected tensor with at least 1 element"); ET_SWITCH_REAL_TYPES_AND( Bool, self_tensor.scalar_type(),