Skip to content

Commit 01263d4

Browse files
committed
Update comment and if condition for input tensor memory alloc_perference
1 parent 0be6f57 commit 01263d4

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
@@ -2004,7 +2004,10 @@ ModelInstanceState::SetInputTensors(
20042004

20052005
// The input must be in contiguous CPU/GPU memory.
20062006
std::vector<std::pair<TRITONSERVER_MemoryType, int64_t>> alloc_perference;
2007-
if (device_.is_cpu()) {
2007+
// For 'KIND_MODEL', intput will always be in CPU as we don't have a way to
2008+
// query the input types.
2009+
if ((device_.is_cpu()) ||
2010+
(Kind() == TRITONSERVER_INSTANCEGROUPKIND_MODEL)) {
20082011
alloc_perference = {{TRITONSERVER_MEMORY_CPU_PINNED, 0},
20092012
{TRITONSERVER_MEMORY_CPU, 0}};
20102013
} else {

0 commit comments

Comments
 (0)