Back out "Move tbe weights as buffers so export can track properly" (pytorch#4469)

Summary:
Pull Request resolved: pytorch#4469

X-link: facebookresearch/FBGEMM#1527

Reviewed By: kqfu

Differential Revision: D78106760

fbshipit-source-id: e8a6e467436e8009b453d7e0f5156268d735ace5

Author: qxy11

fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops_inference.py

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