-
Kate Cheng authored
* Remove if-else and torch.tensor to meet cudagraph requirement Signed-off-by:
Kate Cheng <yunhsuanc@nvidia.com> * Add is_cg_capturable flag to guard the if-else statement Signed-off-by:
Kate Cheng <yunhsuanc@nvidia.com> --------- Signed-off-by:
Kate Cheng <yunhsuanc@nvidia.com> Co-authored-by:
Kirthi Shankar Sivamani <ksivamani@nvidia.com>
aa0659e5