fix tbe reporter (pytorch#4882)

Summary:
Pull Request resolved: pytorch#4882

 - fix mean type
 - fix tensor type

Reviewed By: YanXiong-Meta

Differential Revision: D82248514

fbshipit-source-id: 7ff43e5d3ba8ae36cbda12fa464b595deceb7637

Author: gchalump

fbgemm_gpu/fbgemm_gpu/tbe/bench/tbe_data_config_bench_helper.py

@@ -191,6 +191,7 @@ def generate_requests(
 
     # Generate indices
     all_indices = _generate_indices(tbe_data_config, iters, Bs, L_offsets)
+    all_indices = all_indices.to(get_device())
 
     # Build TBE requests
     if tbe_data_config.variable_B() or tbe_data_config.variable_L():

fbgemm_gpu/fbgemm_gpu/tbe/stats/bench_params_reporter.py

@@ -177,15 +177,19 @@ def extract_params(
         E = (
             Es[0]
             if len(set(Es)) == 1
-            else torch.ceil(torch.mean(torch.tensor(feature_rows)))
+            else torch.ceil(
+                torch.mean(torch.tensor(feature_rows, dtype=torch.float))
+            ).item()
         )
         # Set mixed_dim to be True if there are multiple dims
         mixed_dim = len(set(Ds)) > 1
         # Set D to be the mean of the dims to avoid biasing
         D = (
             Ds[0]
             if not mixed_dim
-            else torch.ceil(torch.mean(torch.tensor(feature_dims)))
+            else torch.ceil(
+                torch.mean(torch.tensor(feature_dims, dtype=torch.float))
+            ).item()
         )
 
         # Compute indices distribution parameters
@@ -198,7 +202,7 @@ def extract_params(
 
         # Compute batch parameters
         batch_params = BatchParams(
-            B=((offsets.numel() - 1) // T),
+            B=int((offsets.numel() - 1) // T),
             sigma_B=(
                 int(
                     torch.ceil(