Skip to content
This repository was archived by the owner on Nov 17, 2023. It is now read-only.

Commit cfbcfb1

Browse files
waytrue17Wei Chu
andauthored
revert cleanup stream (#20998)
Co-authored-by: Wei Chu <weichu@amazon.com>
1 parent 32b65a2 commit cfbcfb1

File tree

2 files changed

+7
-0
lines changed

2 files changed

+7
-0
lines changed

src/engine/naive_engine.cc

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -73,11 +73,14 @@ class NaiveEngine final : public Engine {
7373
LOG(INFO) << "Engine shutdown";
7474
for (size_t i = 0; i < streams_.size(); ++i) {
7575
if (streams_[i] != nullptr) {
76+
// Catch exception for CUDA driver shutdown
77+
MSHADOW_CATCH_ERROR(mshadow::DeleteStream(streams_[i]));
7678
streams_[i] = nullptr;
7779
}
7880
}
7981
for (size_t i = 0; i < aux_streams_.size(); ++i) {
8082
if (aux_streams_[i] != nullptr) {
83+
delete aux_streams_[i];
8184
aux_streams_[i] = nullptr;
8285
}
8386
}

src/engine/threaded_engine_perdevice.cc

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -270,6 +270,10 @@ class ThreadedEnginePerDevice : public ThreadedEngine {
270270
while (task_queue->Pop(&opr_block)) {
271271
this->ExecuteOprBlock(run_ctx, opr_block);
272272
}
273+
// Catch exception for CUDA driver shutdown
274+
MSHADOW_CATCH_ERROR(mshadow::DeleteStream<gpu>(stream));
275+
if (aux_stream != nullptr)
276+
delete aux_stream;
273277
#else
274278
ready_event->signal();
275279
#endif

0 commit comments

Comments
 (0)