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