nit: Resolve lint errors

Signed-off-by: Kshitij Janardan Lakhani <[email protected]>

Author: KshitijLakhani

transformer_engine/pytorch/attention.py

@@ -17,7 +17,6 @@
 from packaging.version import Version as PkgVersion
 
 import torch
-import torch.nn.functional as F
 
 import transformer_engine_torch as tex
 from transformer_engine.pytorch.utils import (

transformer_engine/pytorch/dot_product_attention/inference.py

@@ -2,6 +2,10 @@
 #
 # See LICENSE for license information.
 
+"""
+Inference classes for attention
+"""
+
 
 class InferenceParams:  # pylint: disable=too-few-public-methods
     """

transformer_engine/pytorch/dot_product_attention/rope.py

@@ -2,6 +2,9 @@
 #
 # See LICENSE for license information.
 
+"""
+Rotary Position Embedding implementation of different types along with hlper functions
+"""
 from typing import Optional, Tuple, Union
 import torch
 import transformer_engine_torch as tex

transformer_engine/pytorch/dot_product_attention/utils.py

@@ -2,6 +2,9 @@
 #
 # See LICENSE for license information.
 
+"""
+Utils/Helper classes and methods for attention
+"""
 import math
 import os
 from typing import Any, Dict, List, Optional, Tuple, Union, TypeAlias
@@ -60,7 +63,6 @@ class AttentionLogging:
     _log_level = _NVTE_DEBUG * _NVTE_DEBUG_LEVEL
     _formatter = logging.Formatter("[%(levelname)-8s | %(name)-19s]: %(message)s")
     _stream_handler = logging.StreamHandler()
-    # TODO: Move fa_logger to FAUtils
     fa_logger = logging.getLogger(__name__)
 
     @staticmethod
@@ -1348,7 +1350,10 @@ def _unpack_3_tensors(
 
 
 class PackTensors(torch.autograd.Function):
-
+    """
+    Autograd function to pack a tensor.
+    """
+    
     @staticmethod
     def forward(
         ctx, indices: torch.Tensor, *tensors: Tuple[torch.Tensor, ...]