     else
         return a;
 }
+
+#if __CUDACC_VER_MAJOR__ >= 9
+// specialization for half for handling NaNs
+template <>
+__host__ __device__ __forceinline__ KeyValuePair<int, __half> ArgMin::operator()(
+    const KeyValuePair<int, __half> &a,
+    const KeyValuePair<int, __half> &b) const
+{
+    if (isnan(a.value))
+        return a;
+    else if (isnan(b.value))
+        return b;
+    else if ((b.value < a.value) || ((a.value == b.value) && (b.key < a.key)))
+        return b;
+    else
+        return a;
+}
+#endif
 /* ------------------------------------ End of "patches" ------------------------------------ */
 
 

        return a;
}

//
// ArgMin()
//

/* ------------------------------------ "Patches" to CUB ------------------------------------
   These stubs are needed because CUB does not handle NaNs properly, while NumPy has certain
   behaviors with which we must comply.
*/

//
