Skip to content

Commit

Permalink
explicitly disable pt2 compile on gauc (#2703)
Browse files Browse the repository at this point in the history
Summary:
Pull Request resolved: #2703

Context: the recent changes to GAUC have made it incompatible with PT2 compile, thus causing several issues with launching jobs on the latest trunk. After some digging, it turns out that the incompatibility lies with `torch.Tensor.item()`: pytorch/pytorch#130917.

This diff: explicitly disables PT2 compilation on the problematic function, thus preventing training jobs from crashing due to incompatibility with PT2 compile.

Reviewed By: shz117

Differential Revision: D68629159

fbshipit-source-id: 7cac0c6ebe720f02033ac807a47315fb20996595
  • Loading branch information
Zheng-Yong (Arsa) Ang authored and facebook-github-bot committed Jan 25, 2025
1 parent 4d7b7ff commit 526902f
Showing 1 changed file with 1 addition and 0 deletions.
1 change: 1 addition & 0 deletions torchrec/metrics/gauc.py
Original file line number Diff line number Diff line change
Expand Up @@ -100,6 +100,7 @@ def to_3d(
return torch.ops.fbgemm.jagged_2d_to_dense(tensor_2d, offsets, max_length)


@torch.compiler.disable
def get_auc_states(
labels: torch.Tensor,
predictions: torch.Tensor,
Expand Down

0 comments on commit 526902f

Please sign in to comment.