Skip to content

Commit b389635

Browse files
authored
Fix torch prelude empty-tensor check for zero-numel shapes (#10401)
1 parent 61cd228 commit b389635

File tree

1 file changed

+5
-6
lines changed

1 file changed

+5
-6
lines changed

prelude/slang-torch-prelude.h

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -158,20 +158,19 @@ TensorView make_tensor_view(
158158
.append(")")
159159
.c_str());
160160

161-
bool isEmpty = true;
161+
// A tensor can have zero elements even if some dimensions are non-zero
162+
// (e.g. shape (10, 0)). Emptiness must be based on numel().
163+
bool isEmpty = (val.numel() == 0);
162164
for (int i = 0; i < val.dim(); ++i)
163165
{
166+
res.sizes[i] = val.size(i);
164167
res.strides[i] = val.stride(i) * elementSize;
165-
if (res.strides[i] == 0)
168+
if (!isEmpty && res.strides[i] == 0)
166169
throw std::runtime_error(
167170
std::string(name)
168171
.append(": tensors with broadcasted dimensions are not supported (use "
169172
"tensor.contiguous() to make tensor whole)")
170173
.c_str());
171-
172-
res.sizes[i] = val.size(i);
173-
if (res.sizes[i] > 0)
174-
isEmpty = false;
175174
}
176175

177176
if (!res.data && !isEmpty)

0 commit comments

Comments
 (0)