faster network test

marius-team · Nov 22, 2023 · 8251fe7 · 8251fe7
1 parent eb516eb
commit 8251fe7
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/src/cpp/src/common/util.cpp b/src/cpp/src/common/util.cpp
@@ -135,7 +135,7 @@ torch::Tensor transfer_tensor(torch::Tensor input, torch::Device device, CudaStr
 #ifdef MARIUS_CUDA
         if (device.is_cuda() || input.device().is_cuda()) {
             if (compute_stream != nullptr) input.record_stream(*compute_stream);
-            if (transfer_stream != nullptr) input.record_stream(*transfer_stream);
+//            if (transfer_stream != nullptr) input.record_stream(*transfer_stream);
         }
 #endif
     }

diff --git a/src/cpp/src/pipeline/pipeline_gpu.cpp b/src/cpp/src/pipeline/pipeline_gpu.cpp
@@ -14,7 +14,7 @@
 
 void batchToDevice(Pipeline* pipeline_, shared_ptr<Batch> batch) {
     if (batch->sub_batches_.size() > 0) {
-//        #pragma omp parallel for
+        #pragma omp parallel for
         for (int i = 0; i < batch->sub_batches_.size(); i++) {
             batch->sub_batches_[i]->to(pipeline_->model_->device_models_[i]->device_, pipeline_->dataloader_->compute_streams_[i]);
         }