Skip to content

Commit 467b8d1

Browse files
committed
Update comment and if condition for input tensor memory alloc_perference
1 parent 6f786a7 commit 467b8d1

File tree

1 file changed

+4
-1
lines changed

1 file changed

+4
-1
lines changed

src/libtorch.cc

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1995,7 +1995,10 @@ ModelInstanceState::SetInputTensors(
19951995

19961996
// The input must be in contiguous CPU/GPU memory.
19971997
std::vector<std::pair<TRITONSERVER_MemoryType, int64_t>> alloc_perference;
1998-
if (device_.is_cpu()) {
1998+
// For 'KIND_MODEL', intput will always be in CPU as we don't have a way to
1999+
// query the input types.
2000+
if ((device_.is_cpu()) ||
2001+
(Kind() == TRITONSERVER_INSTANCEGROUPKIND_MODEL)) {
19992002
alloc_perference = {{TRITONSERVER_MEMORY_CPU_PINNED, 0},
20002003
{TRITONSERVER_MEMORY_CPU, 0}};
20012004
} else {

0 commit comments

Comments
 (0)