chore: set stream when new stream is acquired

pytorch · Sep 27, 2024 · 691cf1e · 691cf1e
1 parent 43eb560
commit 691cf1e
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/core/runtime/execute_engine.cpp b/core/runtime/execute_engine.cpp
@@ -305,6 +305,7 @@ std::vector<at::Tensor> execute_engine(std::vector<at::Tensor> inputs, c10::intr
   if (compiled_engine->engine_stream == c10::cuda::getDefaultCUDAStream(current_device_id)) {
     // Create a new stream if the engine stream is the default stream
     compiled_engine->engine_stream = c10::cuda::getStreamFromPool(false, current_device_id);
+    c10::cuda::setCurrentCUDAStream(compiled_engine->engine_stream);
   } else {
     compiled_engine->engine_stream = compiled_engine->caller_stream;
   }

diff --git a/py/torch_tensorrt/dynamo/runtime/_PythonTorchTensorRTModule.py b/py/torch_tensorrt/dynamo/runtime/_PythonTorchTensorRTModule.py
@@ -5,6 +5,7 @@
 from tempfile import tempdir
 from typing import Any, Dict, List, Optional, Sequence, Tuple
 
+import tensorrt as trt
 import torch
 import torch_tensorrt
 from torch.nn import Module
@@ -19,8 +20,6 @@
     multi_gpu_device_check,
 )
 
-import tensorrt as trt
-
 logger = logging.getLogger(__name__)
 
 
@@ -372,6 +371,7 @@ def forward(self, *inputs: torch.Tensor) -> torch.Tensor | Tuple[torch.Tensor, .
                     or self._engine_stream is None
                 ):
                     self._engine_stream = torch.cuda.Stream()
+                    torch.cuda.set_stream(self._engine_stream)
                 else:
                     self._engine_stream = self._caller_stream