diff --git a/torchrec/hybrid_torchrec/src/ids_process/ids_mapper.cpp b/torchrec/hybrid_torchrec/src/ids_process/ids_mapper.cpp
index ceae77d5787eb4ae3e7d1d2b12aac1eb407ede43..9505152c30a7b8709a5e9d9d4a765f868ffbfc51 100644
--- a/torchrec/hybrid_torchrec/src/ids_process/ids_mapper.cpp
+++ b/torchrec/hybrid_torchrec/src/ids_process/ids_mapper.cpp
@@ -1,3 +1,11 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * Copyright (c) Huawei Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
 #include "ids_mapper.h"
 
 #include <algorithm>
diff --git a/torchrec/hybrid_torchrec/src/ids_process/ids_mapper.h b/torchrec/hybrid_torchrec/src/ids_process/ids_mapper.h
index 16de28c6b64ba3446baf1d62a4ae180726160e68..42d3a4e9a16d206ad3a1a342e8796fa9a5508a34 100644
--- a/torchrec/hybrid_torchrec/src/ids_process/ids_mapper.h
+++ b/torchrec/hybrid_torchrec/src/ids_process/ids_mapper.h
@@ -1,3 +1,11 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * Copyright (c) Huawei Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
 #ifndef IDS_MAPPER_H
 #define IDS_MAPPER_H
 
diff --git a/torchrec/hybrid_torchrec/src/ids_process/unique.cpp b/torchrec/hybrid_torchrec/src/ids_process/unique.cpp
index b5b221945362471f079aacb8948580afabde75d4..ce786f21a29f7c09f7fc2b430a9fba4ad6536ce9 100644
--- a/torchrec/hybrid_torchrec/src/ids_process/unique.cpp
+++ b/torchrec/hybrid_torchrec/src/ids_process/unique.cpp
@@ -1,3 +1,12 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * Copyright (c) Huawei Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
 #include "unique.h"
 
 #include <ATen/TensorIndexing.h>
@@ -21,24 +30,12 @@ struct IsUnique {};
 
 template <typename scalar_t>
 struct IsUnique<scalar_t, false> {
-    inline bool operator()(scalar_t* dataPtr, int64_t i)
-    {
-        if (i == 0) {
-            return true;
-        }
-        return c10::load(&dataPtr[i]) != c10::load(&dataPtr[i - 1]);
-    }
-};
-
-template <typename scalar_t>
-struct IsUnique<scalar_t, true> {
-    inline bool operator()(scalar_t* dataPtr, int64_t i)
+    inline bool operator()(scalar_t* sortedDataPtr, int64_t idx)
     {
-        if (i == 0) {
+        if (idx == 0) {
             return true;
         }
-        return (c10::load(&dataPtr[i]) != c10::load(&dataPtr[i - 1])) &&
-               !(_isnan(dataPtr[i]) && _isnan(dataPtr[i - 1]));
+        return c10::load(&sortedDataPtr[idx]) != c10::load(&sortedDataPtr[idx - 1]);
     }
 };
 
diff --git a/torchrec/hybrid_torchrec/src/ids_process/unique.h b/torchrec/hybrid_torchrec/src/ids_process/unique.h
index 58cf7c924254c4a7b26b76e3f4b4eecadf929cdd..74d53cb4acdb04c2d65fd31f6e4f08bc7fe4004b 100644
--- a/torchrec/hybrid_torchrec/src/ids_process/unique.h
+++ b/torchrec/hybrid_torchrec/src/ids_process/unique.h
@@ -1,3 +1,11 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * Copyright (c) Huawei Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
 #ifndef HYBRID_UNIQUE_H
 #define HYBRID_UNIQUE_H
 #include <omp.h>