Fixed movement of tensors across devices

tohtana · Sep 7, 2021 · 4b44949 · 4b44949
1 parent 5afb3c0
commit 4b44949
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 4 deletions.
diff --git a/src/comm/MPIUtil.cpp b/src/comm/MPIUtil.cpp
@@ -6,15 +6,13 @@
 #include <torch/torch.h>
 
 #include "Common.h"
-#include "graph/ir.h"
 #include "MPIUtil.h"
 
 namespace {
 const int RANNC_TAG_UB = 30000;
 }
 
 namespace mpi {
-auto logger = spdlog::stdout_color_mt("mpi");
 
 int getRank() {
   int rank;
@@ -59,7 +57,7 @@ int getTagUB() {
   //// The above returns different values in different ranks.
 
   return RANNC_TAG_UB;
-};
+}
 
 void checkMPIResult(int code) {
   if (code == MPI_SUCCESS)

diff --git a/src/comp/GraphConnector.cpp b/src/comp/GraphConnector.cpp
@@ -432,7 +432,7 @@ std::unordered_map<std::string, IValueMap> GraphConnector::forward(
         // Make sure that copy to cpu has finished
         copy_to_cpu_events_[id][first_split].block(stream);
 
-        inputs_[id][0] =
+        inputs_[id][first_split] =
             toCUDAIfAvailable(inputs_[id][first_split], true, true);
         copy_to_gpu_events_[id][first_split].record(stream);
       }