From e34d84ac2160e14823f54162ff2608b96c97dafc Mon Sep 17 00:00:00 2001
From: Jianyu Huang <jianyuhuang@fb.com>
Date: Thu, 5 May 2022 16:42:49 -0700
Subject: [PATCH] Frontend changes for adding per_sample_weights in boundary
 checker (#1069)

Summary:
Pull Request resolved: https://github.com/pytorch/FBGEMM/pull/1069

Follow up for D35768276 (https://github.com/pytorch/FBGEMM/commit/7be1fcb9e3004a00a4f0ffbe9f7af99a8a2d568a): Spin off the frontend changes in a separate Diff to land later to avoid the BC issue.

Reviewed By: RoshanPAN

Differential Revision: D35882262

fbshipit-source-id: 1d58a8aa53a3f9aa1c36d0926e4d9e4274108783
---
 fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops.py b/fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops.py
index 401032b577..c93814bade 100644
--- a/fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops.py
+++ b/fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops.py
@@ -599,6 +599,7 @@ def forward(
                 offsets,
                 self.bounds_check_mode_int,
                 self.bounds_check_warning,
+                per_sample_weights,
             )
         self.step += 1
         if len(self.timesteps_prefetched) == 0:
@@ -1972,6 +1973,7 @@ def forward(
                 offsets,
                 self.bounds_check_mode_int,
                 self.bounds_check_warning,
+                per_sample_weights,
             )
         # Note: CPU and CUDA ops use the same interface to facilitate JIT IR
         # generation for CUDA/CPU. For CPU op, we don't need weights_uvm and