[PyTorch] Fix incorrect docstrings in tensor saving functions (#1549)

Fix incorrect docstrings in tensor saving functions
Signed-off-by: Tim Moon <tmoon@nvidia.com>

transformer_engine/pytorch/tensor/_internal/float8_tensor_base.py

@@ -98,12 +98,7 @@ class Float8TensorBase:
         }
 
     def prepare_for_saving(self) -> Tuple[list[Optional[torch.Tensor]], Float8TensorBase]:
-        """Prepare the tensor base for saving for backward
-
-        After calling this, the tensor instance does not hold any
-        data.
-
-        """
+        """Prepare the tensor base for saving for backward"""
         tensors = [self._data, self._transpose]
         return tensors, self
 

transformer_engine/pytorch/tensor/_internal/mxfp8_tensor_base.py

@@ -93,12 +93,7 @@ class MXFP8TensorBase:
         }
 
     def prepare_for_saving(self) -> Tuple[list[Optional[torch.Tensor]], MXFP8TensorBase]:
-        """Prepare the tensor base for saving for backward
-
-        After calling this, the tensor instance does not hold any
-        data.
-
-        """
+        """Prepare the tensor base for saving for backward"""
         tensors = [self._rowwise_data, self._columnwise_data]
         return tensors, self
 

transformer_engine/pytorch/tensor/float8_tensor.py

@@ -349,12 +349,7 @@ class Float8Tensor(Float8TensorBase, QuantizedTensor):
         self._transpose_invalid = True
 
     def prepare_for_saving(self) -> Tuple[list[Optional[torch.Tensor]], Float8TensorBase]:
-        """Prepare the tensor base for saving for backward
-
-        After calling this, the tensor instance does not hold any
-        data.
-
-        """
+        """Prepare the tensor base for saving for backward"""
         return [self], None
 
     @classmethod

transformer_engine/pytorch/tensor/mxfp8_tensor.py

@@ -286,12 +286,7 @@ class MXFP8Tensor(MXFP8TensorBase, QuantizedTensor):
         self._columnwise_data = torch.Tensor() if self._columnwise_data is not None else None
 
     def prepare_for_saving(self) -> Tuple[list[Optional[torch.Tensor]], MXFP8TensorBase]:
-        """Prepare the tensor base for saving for backward
-
-        After calling this, the tensor instance does not hold any
-        data.
-
-        """
+        """Prepare the tensor base for saving for backward"""
         return [self], None
 
     @classmethod