Skip to content

Commit 0be6f57

Browse files
committed
Add missing #ifdef
1 parent 6a4d9c4 commit 0be6f57

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
@@ -2282,13 +2282,15 @@ void
22822282
ModelInstanceState::SetCurrentCudaStream(
22832283
const cudaStream_t& stream, const int& device_id)
22842284
{
2285+
#ifdef TRITON_ENABLE_GPU
22852286
at::cuda::CUDAStream torch_stream =
22862287
at::cuda::getStreamFromExternal(stream, device_id);
22872288
// This function replaces the default stream with the stream we created. It
22882289
// is not necessary to change the current device to the desired device when
22892290
// replacing the default stream for that device. See the documentation here:
22902291
// https://pytorch.org/cppdocs/api/function_namespacec10_1_1cuda_1a6ed50cc0fc16cc7014d9c2f4c3bd098d.html
22912292
at::cuda::setCurrentCUDAStream(torch_stream);
2293+
#endif
22922294
}
22932295

22942296
/////////////

0 commit comments

Comments
 (0)