Skip to content

Commit 0f2ce85

Browse files
Skip dimension of size 1 in contiguous checks. (triton-inference-server#281)
* Skip dimension of size 1 in contiguous checks. Since PyTorch 1.13 dimension of size 1 have normalised strides in PyTorch which fail here when using DLPack. This was done to conform the torch stride representation and the one from numpy. Unfortunately this means we are stuck with PyTorch 1.12.0 in python models. * Conform to pre-commit guidelines
1 parent cb53f0a commit 0f2ce85

File tree

1 file changed

+7
-5
lines changed

1 file changed

+7
-5
lines changed

src/pb_tensor.cc

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -433,12 +433,14 @@ PbTensor::FromDLPackCapsule(
433433
int64_t calculated_stride{1};
434434
bool is_contiguous_c_order = true;
435435
for (size_t i = 1; i < dims.size(); i++) {
436-
if (strides[ndim - i] != calculated_stride) {
437-
is_contiguous_c_order = false;
438-
break;
439-
}
436+
if (dims[ndim - i] != 1) {
437+
if (strides[ndim - i] != calculated_stride) {
438+
is_contiguous_c_order = false;
439+
break;
440+
}
440441

441-
calculated_stride *= dims[ndim - i];
442+
calculated_stride *= dims[ndim - i];
443+
}
442444
}
443445

444446
if (!is_contiguous_c_order) {

0 commit comments

Comments
 (0)