Skip to content

Commit 37bd812

Browse files
authored
[RUNTIME][CRT] support DLTensor whose ndim == 0 (#5344)
Signed-off-by: windclarion <windclarion@gmail.com>
1 parent 84d1eec commit 37bd812

File tree

1 file changed

+3
-3
lines changed

1 file changed

+3
-3
lines changed

src/runtime/crt/ndarray.c

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -68,15 +68,15 @@ int TVMNDArray_Load(TVMNDArray * ret, const char ** strm) {
6868
ctx = ((DLContext*)*strm)[0]; *strm += sizeof(ctx); // NOLINT(*)
6969
ndim = ((uint32_t*)*strm)[0]; *strm += sizeof(ndim); // NOLINT(*)
7070
dtype = ((DLDataType*)*strm)[0]; *strm += sizeof(dtype); // NOLINT(*)
71-
if ((ndim <= 0) || (ndim > TVM_CRT_MAX_NDIM)) {
72-
fprintf(stderr, "Invalid ndim=%d: expected to be 1 ~ %d.\n", ndim, TVM_CRT_MAX_NDIM);
71+
if ((ndim < 0) || (ndim > TVM_CRT_MAX_NDIM)) {
72+
fprintf(stderr, "Invalid ndim=%d: expected to be 0 ~ %d.\n", ndim, TVM_CRT_MAX_NDIM);
7373
status = -1;
7474
}
7575
if (ctx.device_type != kDLCPU) {
7676
fprintf(stderr, "Invalid DLTensor context: can only save as CPU tensor\n");
7777
status = -1;
7878
}
79-
int64_t shape[TVM_CRT_MAX_NDIM];
79+
int64_t shape[TVM_CRT_MAX_NDIM] = {0};
8080
uint32_t idx;
8181
if (ndim != 0) {
8282
for (idx = 0; idx < ndim; idx++) {

0 commit comments

Comments
 (0)