Move tbe weights as buffers so export can track properly (#4369)

Summary:
X-link: facebookresearch/FBGEMM#1438

Pull Request resolved: #4369

Title

Reviewed By: SherlockNoMad

Differential Revision: D76846172

fbshipit-source-id: 2074df3e5d9d971dba1c9442c8399ae5fc82297e

Author: tugsbayasgalan

fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops_inference.py

@@ -517,18 +517,29 @@ def max_ty_D(ty: SparseType) -> int:
         )
         self.weight_initialized: bool = False
 
-        self.weights_dev: torch.Tensor = torch.zeros(
-            0,
-            device=self.current_device,
-            dtype=torch.uint8,
+        self.register_buffer(
+            "weights_dev",
+            torch.zeros(
+                0,
+                device=self.current_device,
+                dtype=torch.uint8,
+            ),
+            persistent=False,
         )
 
-        self.weights_host: torch.Tensor = torch.zeros(
-            0, device=self.current_device, dtype=torch.uint8
+        self.register_buffer(
+            "weights_host",
+            torch.zeros(
+                0,
+                device=self.current_device,
+                dtype=torch.uint8,
+            ),
         )
 
-        self.weights_uvm: torch.Tensor = torch.empty(
-            0, device=self.current_device, dtype=torch.uint8
+        self.register_buffer(
+            "weights_uvm",
+            torch.empty(0, device=self.current_device, dtype=torch.uint8),
+            persistent=False,
         )
 
         cached_dims = [