Fix invocations of deprecated torch api

fbgemm_gpu/fbgemm_gpu/tbe/bench/eeg_cli.py

@@ -35,7 +35,7 @@ def estimate(indices: str) -> None:
         estimate --indices="indices.pt"
     """
 
-    indices = torch.load(indices)
+    indices = torch.load(indices, weights_only=True)
     heavy_hitters, q, s, max_index, num_indices = (
         torch.ops.fbgemm.tbe_estimate_indices_distribution(indices)
     )

fbgemm_gpu/fbgemm_gpu/tbe/cache/kv_embedding_ops_inference.py

@@ -270,7 +270,7 @@ def fill_random_weights(self) -> None:
             embedding_dim = rounded_row_size_in_bytes(
                 embedding_dim, weight_ty, self.row_alignment
             )
-            indices = torch.range(0, num_embeddings - 1, dtype=torch.int64)
+            indices = torch.arange(num_embeddings, dtype=torch.int64)
             weights = random_quant_scaled_tensor(
                 shape=torch.Size([num_embeddings, embedding_dim]),
                 device=self.current_device,

fbgemm_gpu/fbgemm_gpu/tbe/utils/requests.py

@@ -77,13 +77,15 @@ def generate_requests_from_data_file(
     ), "If requests_data_file is provided, indices_file and offsets_file cannot be provided."
 
     if requests_data_file:
-        indices_tensor, offsets_tensor, *rest = torch.load(requests_data_file)
+        indices_tensor, offsets_tensor, *rest = torch.load(
+            requests_data_file, weights_only=True
+        )
     else:
         assert (
             indices_file and offsets_file
         ), "Both indices_file and offsets_file must be provided if either is provided."
-        indices_tensor = torch.load(indices_file)
-        offsets_tensor = torch.load(offsets_file)
+        indices_tensor = torch.load(indices_file, weights_only=True)
+        offsets_tensor = torch.load(offsets_file, weights_only=True)
 
     average_L = 0
     if tables is not None:

fbgemm_gpu/test/quantize/fp8_rowwise_test.py

@@ -225,7 +225,7 @@ def test_quantize_and_dequantize_op_padded_fp8_rowwise(
                 f"mean: {errors.abs().mean() * 100:.1f}%"
             )
 
-        torch.testing.assert_allclose(dqcat, qref, rtol=0.1, atol=0.05)
+        torch.testing.assert_close(dqcat, qref, rtol=0.1, atol=0.05)
 
     @unittest.skipIf(*gpu_unavailable)
     def test_padded_fp8_rowwise_input_validation(self) -> None:

fbgemm_gpu/test/sparse/block_bucketize_test.py

@@ -1186,9 +1186,11 @@ def test_block_bucketize_sparse_features_inference(
             torch.testing.assert_close(
                 new_lengths_gpu.cpu(), new_lengths_ref, rtol=0, atol=0
             )
-            torch.testing.assert_allclose(
+            torch.testing.assert_close(
                 bucket_mapping_gpu.cpu(),
                 bucket_mapping,
+                rtol=0,
+                atol=0,
             )
 
     @given(
@@ -1289,9 +1291,11 @@ def test_populate_bucketized_permute(
             torch.testing.assert_close(
                 new_lengths_gpu.cpu(), new_lengths_ref, rtol=0, atol=0
             )
-            torch.testing.assert_allclose(
+            torch.testing.assert_close(
                 bucket_mapping_gpu.cpu(),
                 bucket_mapping_cpu,
+                rtol=0,
+                atol=0,
             )
 
     def _run_populate_bucketized_permute_test(

fbgemm_gpu/test/utils/filestore_test.py

@@ -46,7 +46,9 @@ def _test_filestore_readwrite(
         assert store.exists(path), f"{path} does not exist"
 
         if isinstance(input, torch.Tensor):
-            assert torch.load(store.read(path)).equal(input), "tensors do not match"
+            assert torch.load(store.read(path), weights_only=True).equal(
+                input
+            ), "tensors do not match"
 
         elif isinstance(input, io.BytesIO) or isinstance(input, BinaryIO):
             input.seek(0)