fix bug that causes servicelab issue (#3605)

Summary:
Pull Request resolved: #3605

# context
* previous change in `GPUMemoryStats` is causing issue in servicelabe
* fix the bug
> We found a Tupperware task failure only on the control side

Reviewed By: spmex

Differential Revision: D88860603

fbshipit-source-id: 4d65b232c5b1cd3da5024f5a0d74ad5743ea17db

Author: TroyGarden

torchrec/distributed/benchmark/base.py

@@ -504,14 +504,6 @@ def wrapper() -> Any:  # pyre-ignore [3]
                 help="JSON config file for benchmarking",
             )
 
-            # Add loglevel argument with current logger level as default
-            parser.add_argument(
-                "--loglevel",
-                type=str,
-                default=logging._levelToName[logger.level],
-                help="Set the logging level (e.g. info, debug, warning, error)",
-            )
-
             pre_args, _ = parser.parse_known_args()
 
             yaml_defaults: Dict[str, Any] = (
@@ -531,7 +523,6 @@ def wrapper() -> Any:  # pyre-ignore [3]
             seen_args = {
                 "json_config",
                 "yaml_config",
-                "loglevel",
             }
 
             for _name, param in sig.parameters.items():

torchrec/sparse/tests/jagged_tensor_benchmark.py

@@ -110,7 +110,7 @@ def wrapped_func(
             short_name=name,
             gpu_elapsed_time=torch.tensor(times) * 1e3,
             cpu_elapsed_time=torch.tensor(times) * 1e3,
-            gpu_mem_stats=[GPUMemoryStats(0, 0, 0, 0)],
+            gpu_mem_stats=[GPUMemoryStats(0, 0, 0, 0, 0, 0)],
             cpu_mem_stats=[CPUMemoryStats.for_process(0)],
         )
 

torchrec/sparse/tests/keyed_jagged_tensor_benchmark_lib.py

@@ -233,7 +233,7 @@ def benchmark_kjt(
         short_name=f"{test_name}-{transform_type.name}",
         gpu_elapsed_time=torch.tensor(times),
         cpu_elapsed_time=torch.tensor(times),
-        gpu_mem_stats=[GPUMemoryStats(0, 0, 0, 0)],
+        gpu_mem_stats=[GPUMemoryStats(0, 0, 0, 0, 0, 0)],
         cpu_mem_stats=[CPUMemoryStats.for_process(0)],
     )