add ref

Signed-off-by: Youngeun Kwon <[email protected]>
NVIDIA · Dec 5, 2024 · 689e30a · 689e30a
1 parent daba5a6
commit 689e30a
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py
@@ -447,6 +447,7 @@ def fsdp_pre_all_gather(self, mesh):  # pylint: disable=unused-argument
         """
         A hook function used in torch fsdp2, called before all-gather
         return (all-gather input), (metadata)
+        Ref: https://github.com/pytorch/pytorch/pull/122908
 
         """
 
@@ -463,6 +464,7 @@ def fsdp_post_all_gather(
         """
         A hook function used in torch fsdp2, called after all-gather
         return (Float8Tensor class instance of all-gathered input), (Things to free after forward)
+        Ref: https://github.com/pytorch/pytorch/pull/122908
 
         """
         (data,) = all_gather_outputs