Use proper type for pooling_mode (pytorch#1211)

Summary: Pull Request resolved: pytorch#1211 As title. Reviewed By: jianyuh Differential Revision: D37877585 fbshipit-source-id: 8063c05c5ed03b3c9c76c3a20bfb41db332a6d19
williamwen42 · Jul 20, 2022 · 58a3df0 · 58a3df0
1 parent c9bbb77
commit 58a3df0
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops.py b/fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops.py
@@ -2064,7 +2064,7 @@ def forward(
             max_float32_D=self.max_float32_D,
             indices=indices,
             offsets=offsets,
-            pooling_mode=self.pooling_mode,
+            pooling_mode=int(self.pooling_mode),
             indice_weights=per_sample_weights,
             output_dtype=self.output_dtype,
             lxu_cache_weights=self.lxu_cache_weights,