Skip to content

Commit 6f786a7

Browse files
committed
Add missing #ifdef
1 parent 7775afb commit 6f786a7

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
@@ -2272,13 +2272,15 @@ void
22722272
ModelInstanceState::SetCurrentCudaStream(
22732273
const cudaStream_t& stream, const int& device_id)
22742274
{
2275+
#ifdef TRITON_ENABLE_GPU
22752276
at::cuda::CUDAStream torch_stream =
22762277
at::cuda::getStreamFromExternal(stream, device_id);
22772278
// This function replaces the default stream with the stream we created. It
22782279
// is not necessary to change the current device to the desired device when
22792280
// replacing the default stream for that device. See the documentation here:
22802281
// https://pytorch.org/cppdocs/api/function_namespacec10_1_1cuda_1a6ed50cc0fc16cc7014d9c2f4c3bd098d.html
22812282
at::cuda::setCurrentCUDAStream(torch_stream);
2283+
#endif
22822284
}
22832285

22842286
/////////////

0 commit comments

Comments
 (0)