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 @@ -2282,13 +2282,15 @@ void
2282
2282
ModelInstanceState::SetCurrentCudaStream (
2283
2283
const cudaStream_t& stream, const int & device_id)
2284
2284
{
2285
+ #ifdef TRITON_ENABLE_GPU
2285
2286
at::cuda::CUDAStream torch_stream =
2286
2287
at::cuda::getStreamFromExternal (stream, device_id);
2287
2288
// This function replaces the default stream with the stream we created. It
2288
2289
// is not necessary to change the current device to the desired device when
2289
2290
// replacing the default stream for that device. See the documentation here:
2290
2291
// https://pytorch.org/cppdocs/api/function_namespacec10_1_1cuda_1a6ed50cc0fc16cc7014d9c2f4c3bd098d.html
2291
2292
at::cuda::setCurrentCUDAStream (torch_stream);
2293
+ #endif
2292
2294
}
2293
2295
2294
2296
// ///////////
You can’t perform that action at this time.
0 commit comments