From bd8d904d71d1a5597bf3fd3fc795945a1a762c0a Mon Sep 17 00:00:00 2001 From: Yan Xiong Date: Fri, 5 Sep 2025 14:00:15 -0700 Subject: [PATCH] convert batch size to float before torch.std in params reporter Summary: ### Description This diff converts the batch size to a float before calculating the standard deviation in the `params_reporter` module. This change is made to ensure accurate calculations, as `torch.std` expects a floating-point input. ### Changes - **Modified code in `bench_params_reporter.py`** - Changed the line `torch.std(torch.tensor([b for bs in batch_size_per_feature_per_rank for b in bs]))` to `torch.std(torch.tensor([b for bs in batch_size_per_feature_per_rank for b in bs])).float()`. ### Reason The `torch.std` function requires a floating-point tensor to calculate the standard deviation. By converting the batch size to a float, we ensure that the calculation is performed correctly. Differential Revision: D81809491 --- fbgemm_gpu/fbgemm_gpu/tbe/stats/bench_params_reporter.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/fbgemm_gpu/fbgemm_gpu/tbe/stats/bench_params_reporter.py b/fbgemm_gpu/fbgemm_gpu/tbe/stats/bench_params_reporter.py index 5df112257e..9dc3372ce1 100644 --- a/fbgemm_gpu/fbgemm_gpu/tbe/stats/bench_params_reporter.py +++ b/fbgemm_gpu/fbgemm_gpu/tbe/stats/bench_params_reporter.py @@ -209,7 +209,7 @@ def extract_params( for bs in batch_size_per_feature_per_rank for b in bs ] - ) + ).float() ) ) )