Add logger warning stacklevel (#3242)

aliafzal · facebook-github-bot · commit 2e874f7515be · 2025-07-30T20:06:48.000-07:00
Summary: Pull Request resolved: #3242 TSIA Reviewed By: SSYernar Differential Revision: D79149808 fbshipit-source-id: 3fc79db614571ad556798f70581c6e40c7b738f9
diff --git a/torchrec/distributed/sharding_plan.py b/torchrec/distributed/sharding_plan.py
@@ -257,6 +257,7 @@ def _get_block_size_for_cw_shard(
             f"Dim of {columns} cannot be evenly divided with column wise shard"
             "dim {column_wise_shard_dim}, overriding block_size to embedding_dim={columns}",
             UserWarning,
+            stacklevel=2,
         )
         block_size = columns
     return block_size
@@ -281,6 +282,7 @@ def _calculate_cw_shard_sizes_and_offsets(
             f"Dim of {columns} cannot be evenly divided with column wise shard"
             "dim {col_wise_shard_dim}, overriding block_size to embedding_dim={columns}",
             UserWarning,
+            stacklevel=2,
         )
         block_size = columns
 

Original file line number	Diff line number	Diff line change
`@@ -257,6 +257,7 @@ def _get_block_size_for_cw_shard(`
`257`	`257`	`f"Dim of {columns} cannot be evenly divided with column wise shard"`
`258`	`258`	`"dim {column_wise_shard_dim}, overriding block_size to embedding_dim={columns}",`
`259`	`259`	`UserWarning,`
	`260`	`+ stacklevel=2,`
`260`	`261`	`)`
`261`	`262`	`block_size = columns`
`262`	`263`	`return block_size`
`@@ -281,6 +282,7 @@ def _calculate_cw_shard_sizes_and_offsets(`
`281`	`282`	`f"Dim of {columns} cannot be evenly divided with column wise shard"`
`282`	`283`	`"dim {col_wise_shard_dim}, overriding block_size to embedding_dim={columns}",`
`283`	`284`	`UserWarning,`
	`285`	`+ stacklevel=2,`
`284`	`286`	`)`
`285`	`287`	`block_size = columns`
`286`	`288`