Skip to content

Commit 749901a

Browse files
committed
Add missing #ifdef
1 parent b3d6ec6 commit 749901a

File tree

1 file changed

+2
-0
lines changed

1 file changed

+2
-0
lines changed

src/libtorch.cc

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2207,13 +2207,15 @@ void
22072207
ModelInstanceState::SetCurrentCudaStream(
22082208
const cudaStream_t& stream, const int& device_id)
22092209
{
2210+
#ifdef TRITON_ENABLE_GPU
22102211
at::cuda::CUDAStream torch_stream =
22112212
at::cuda::getStreamFromExternal(stream, device_id);
22122213
// This function replaces the default stream with the stream we created. It
22132214
// is not necessary to change the current device to the desired device when
22142215
// replacing the default stream for that device. See the documentation here:
22152216
// https://pytorch.org/cppdocs/api/function_namespacec10_1_1cuda_1a6ed50cc0fc16cc7014d9c2f4c3bd098d.html
22162217
at::cuda::setCurrentCUDAStream(torch_stream);
2218+
#endif
22172219
}
22182220

22192221
/////////////

0 commit comments

Comments
 (0)