Add support for Rowwise and Columnwise sharding

Differential Revision: D52641307
pytorch · Jan 10, 2024 · 51554a2 · 51554a2
1 parent 9bc6103
commit 51554a2
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 = [