2 files changed, 24 insertions, 5 deletions
diff --git a/Source/Particles/Sorting/Partition.cpp b/Source/Particles/Sorting/Partition.cpp
index 3ee49cafe..d64b9fa25 100644
--- a/Source/Particles/Sorting/Partition.cpp
+++ b/Source/Particles/Sorting/Partition.cpp
@@ -50,7 +50,7 @@ PhysicalParticleContainer::PartitionParticlesInBuffers(
     Gpu::ManagedDeviceVector<long> pid;
     pid.resize(np);
 
-    // First, partition particles in the larger buffer
+    // First, partition particles into the larger buffer
 
     // - Select the larger buffer
     iMultiFab const* bmasks =
@@ -68,16 +68,22 @@ PhysicalParticleContainer::PartitionParticlesInBuffers(
     // separates the particles that deposit/gather on the fine patch (first part)
     // and the particles that deposit/gather in the buffers (last part)
 
+    // Second, among particles that are in the larger buffer, partition
+    // particles into the smaller buffer
     if (WarpX::n_current_deposition_buffer == WarpX::n_field_gather_buffer) {
-        nfine_current = nfine_gather = std::distance(pid.begin(), sep);
-    } else if (sep != pid.end()) {
+        // No need to do anything if the buffers have the same size
+        nfine_current = nfine_gather = iteratorDistance(pid.begin(), sep);
+    } else if (sep == pid.end()) {
+        // No need to do anything if there are no particles in the larger buffer
+        nfine_current = nfine_gather = np;
+    } else {
         int n_buf;
         if (bmasks == gather_masks) {
-            nfine_gather = std::distance(pid.begin(), sep);
+            nfine_gather = iteratorDistance(pid.begin(), sep);
             bmasks = current_masks;
             n_buf = WarpX::n_current_deposition_buffer;
         } else {
-            nfine_current = std::distance(pid.begin(), sep);
+            nfine_current = iteratorDistance(pid.begin(), sep);
             bmasks = gather_masks;
             n_buf = WarpX::n_field_gather_buffer;
         }
diff --git a/Source/Particles/Sorting/SortingUtils.H b/Source/Particles/Sorting/SortingUtils.H
index ede59b53b..7d53a352e 100644
--- a/Source/Particles/Sorting/SortingUtils.H
+++ b/Source/Particles/Sorting/SortingUtils.H
@@ -43,6 +43,19 @@ ForwardIterator stablePartition(ForwardIterator index_begin,
 }
 
 // TODO: Add documentation
+template< typename ForwardIterator >
+int iteratorDistance(ForwardIterator first,
+                     ForwardIterator last)
+{
+#ifdef AMREX_USE_GPU
+    // On GPU: Use thrust
+    return thrust::distance( first, last );
+#else
+    return std::distance( first, last );
+#endif
+}
+
+// TODO: Add documentation
 class fillBufferFlag
 {
     public: