From 21c4ffe75217b3b314cc4dbd3c2e71f98c498129 Mon Sep 17 00:00:00 2001
From: Adrian-Diaz <adriandiaz1117@gmail.com>
Date: Mon, 16 Dec 2024 15:07:33 -0700
Subject: [PATCH] BUG: let it sync as it needs to

---
 src/include/tpetra_wrapper_types.h | 16 ++++++----------
 1 file changed, 6 insertions(+), 10 deletions(-)

diff --git a/src/include/tpetra_wrapper_types.h b/src/include/tpetra_wrapper_types.h
index 7c30636..53d36e9 100644
--- a/src/include/tpetra_wrapper_types.h
+++ b/src/include/tpetra_wrapper_types.h
@@ -1520,13 +1520,11 @@ void TpetraDCArray<T,Layout,ExecSpace,MemoryTraits>::repartition_vector() {
     
     //for whatever reason, when using one process the device contains the updated data, when using several the host does
     //so we need this if block
-    if(nranks==1){
-        this_array_.template modify<typename TArray1D::execution_space>();
-        this_array_.template sync<typename TArray1D::host_mirror_space>();
+    if(this_array_.template need_sync<typename TArray1D::execution_space>()){
+        this_array_.template sync<typename TArray1D::execution_space>();
     }
     else{
-        this_array_.template modify<typename TArray1D::host_mirror_space>();
-        this_array_.template sync<typename TArray1D::execution_space>();
+        this_array_.template sync<typename TArray1D::host_mirror_space>();
     }
 }
 
@@ -2755,13 +2753,11 @@ void TpetraDFArray<T,Layout,ExecSpace,MemoryTraits>::repartition_vector() {
 
     //for whatever reason, when using one process the device contains the updated data, when using several the host does
     //so we need this if block
-    if(nranks==1){
-        this_array_.template modify<typename TArray1D::execution_space>();
-        this_array_.template sync<typename TArray1D::host_mirror_space>();
+    if(this_array_.template need_sync<typename TArray1D::execution_space>()){
+        this_array_.template sync<typename TArray1D::execution_space>();
     }
     else{
-        this_array_.template modify<typename TArray1D::host_mirror_space>();
-        this_array_.template sync<typename TArray1D::execution_space>();
+        this_array_.template sync<typename TArray1D::host_mirror_space>();
     }
 
 }