Make initialize_weights scriptable in IntNBitTableBatchedEmbeddingBagsCodegen (#1622)

Summary:
Pull Request resolved: #1622

Need to init the weight during model loading with XL V2, so export it.

Reviewed By: qxy11, jianyuh

Differential Revision: D43709984

fbshipit-source-id: 1730ea6f108303fee0811eb9de3fa5bef2fa8267

Author: houseroad

fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops.py

@@ -2451,7 +2451,8 @@ def _apply_split(
         if uvm_size > 0:
             assert not self.use_cpu
             if enforce_hbm:
-                logging.info("Enforce hbm for the cache location")
+                if not torch.jit.is_scripting():
+                    logging.info("Enforce hbm for the cache location")
                 self.weights_uvm = torch.zeros(
                     uvm_size,
                     device=self.current_device,
@@ -2767,6 +2768,7 @@ def split_embedding_weights(
 
         return splits
 
+    @torch.jit.export
     def initialize_weights(self) -> None:
         if not self.weight_initialized:
             self._apply_split(
@@ -2777,7 +2779,7 @@ def initialize_weights(self) -> None:
                 self.weights_physical_offsets,
                 self.enforce_hbm,
             )
-            self.weight_initialized: bool = True
+            self.weight_initialized = True
 
     def fill_random_weights(self) -> None:
         """