From 3172e6f77c5c018f250e67a3013ee31f32e085ed Mon Sep 17 00:00:00 2001
From: Wenjie Yang <wenjieyang@meta.com>
Date: Thu, 10 Oct 2024 03:51:34 -0700
Subject: [PATCH] convert data type of indices and offsets in
 int_nbit_split_embedding_codegen_lookup_function

Summary:
X-link: https://github.com/facebookresearch/FBGEMM/pull/336

Convert offsets.dtype to indices.dtype in gpu ops. The two tensor need to keep same dtype.

Reviewed By: 842974287

Differential Revision: D64085749

fbshipit-source-id: 0c07d85267f11a097d9dd6711fe4591cb5e11226
---
 .../codegen/inference/embedding_forward_quantized_host.cpp     | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/fbgemm_gpu/codegen/inference/embedding_forward_quantized_host.cpp b/fbgemm_gpu/codegen/inference/embedding_forward_quantized_host.cpp
index 8bc20c215..5fcc3a017 100644
--- a/fbgemm_gpu/codegen/inference/embedding_forward_quantized_host.cpp
+++ b/fbgemm_gpu/codegen/inference/embedding_forward_quantized_host.cpp
@@ -283,6 +283,9 @@ Tensor int_nbit_split_embedding_codegen_lookup_function(
     std::optional<int64_t> max_float8_D,
     std::optional<int64_t> fp8_exponent_bits,
     std::optional<int64_t> fp8_exponent_bias) {
+  if (offsets.scalar_type() != indices.scalar_type()) {
+    offsets = offsets.toType(indices.scalar_type());
+  }
   if (static_cast<PoolingMode>(pooling_mode) == PoolingMode::NONE) {
     std::vector<int64_t> max_D_list{
         max_int2_D,