File tree Expand file tree Collapse file tree 1 file changed +2
-0
lines changed Expand file tree Collapse file tree 1 file changed +2
-0
lines changed Original file line number Diff line number Diff line change @@ -2207,13 +2207,15 @@ void
2207
2207
ModelInstanceState::SetCurrentCudaStream (
2208
2208
const cudaStream_t& stream, const int & device_id)
2209
2209
{
2210
+ #ifdef TRITON_ENABLE_GPU
2210
2211
at::cuda::CUDAStream torch_stream =
2211
2212
at::cuda::getStreamFromExternal (stream, device_id);
2212
2213
// This function replaces the default stream with the stream we created. It
2213
2214
// is not necessary to change the current device to the desired device when
2214
2215
// replacing the default stream for that device. See the documentation here:
2215
2216
// https://pytorch.org/cppdocs/api/function_namespacec10_1_1cuda_1a6ed50cc0fc16cc7014d9c2f4c3bd098d.html
2216
2217
at::cuda::setCurrentCUDAStream (torch_stream);
2218
+ #endif
2217
2219
}
2218
2220
2219
2221
// ///////////
You can’t perform that action at this time.
0 commit comments