torchrec
torchrec copied to clipboard
KJT dynamo path for conditions
Summary: The change is no-change for eager path.
Dynamo path
Dynamo tracing can not pass through various conditions => making dynamo path without those conditions.
We skip asserts for dynamo and variable batch check if it can go through "equal batches" logic => for Dynamo fallback on Variable Batch path.
Comments about torch scripting
guard_size_oblivious(), _check_is_size() is not torch scriptable atm, we have to guard it in is not torch.jit.is_scripting
torch.jit._unwrap_optional(recat) is used to hint torch script that recat is not None at the places where we use it.
torch.jit.annotate() hints to TorchScript about the returned type of tensor.tolist()
Differential Revision: D55695198
This pull request was exported from Phabricator. Differential Revision: D55695198
This pull request was exported from Phabricator. Differential Revision: D55695198
This pull request was exported from Phabricator. Differential Revision: D55695198
This pull request was exported from Phabricator. Differential Revision: D55695198