Add support for Rowwise and Columnwise sharding (#1621)

Summary: Add in Rowwise and Columnwise sharding to inference module benchmarking suite by default. Differential Revision: D52641307
pytorch · Jan 10, 2024 · c165e4c · c165e4c
1 parent 9bc6103
commit c165e4c
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/torchrec/distributed/benchmark/benchmark_inference.py b/torchrec/distributed/benchmark/benchmark_inference.py
@@ -439,11 +439,10 @@ def main() -> None:
         # Place all outputs under the datetime folder
         os.mkdir(output_dir)
 
-    # TODO: ROW_WISE and COLUMN_WISE are not supported yet
     BENCH_SHARDING_TYPES = [
         ShardingType.TABLE_WISE,
-        # ShardingType.ROW_WISE,
-        # ShardingType.COLUMN_WISE,
+        ShardingType.ROW_WISE,
+        ShardingType.COLUMN_WISE,
     ]
 
     table_sizes = [