More lint

torchtitan/models/llama/model.py

@@ -8,11 +8,7 @@
 # Copyright (c) Meta Platforms, Inc. All Rights Reserved.
 
 
-from dataclasses import dataclass
-from typing import Optional, Tuple
-
 import torch
-import torch.nn.functional as F
 from torch import nn
 from torchtitan.models.types import ModelArgs
 from torchtitan.modules.attention import Attention

torchtitan/models/types.py

@@ -9,10 +9,8 @@
 
 
 from dataclasses import dataclass
-from typing import Callable, List, Optional, Tuple
+from typing import Callable, List, Optional
 
-import torch
-import torch.nn.functional as F
 from torch import nn
 
 

torchtitan/modules/attention.py

@@ -8,18 +8,11 @@
 # Copyright (c) Meta Platforms, Inc. All Rights Reserved.
 
 
-from dataclasses import dataclass
-from typing import Optional, Tuple
-
 import torch
 import torch.nn.functional as F
 from torch import nn
 from torchtitan.models.types import ModelArgs
-from torchtitan.modules.attention_utils import (
-    apply_rotary_emb,
-    repeat_kv,
-    reshape_for_broadcast,
-)
+from torchtitan.modules.attention_utils import apply_rotary_emb, repeat_kv
 
 
 class Attention(nn.Module):

torchtitan/modules/attention_utils.py

@@ -8,12 +8,9 @@
 # Copyright (c) Meta Platforms, Inc. All Rights Reserved.
 
 
-from dataclasses import dataclass
-from typing import Optional, Tuple
+from typing import Tuple
 
 import torch
-import torch.nn.functional as F
-from torch import nn
 
 
 def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor) -> torch.Tensor:

torchtitan/modules/feed_forward.py

@@ -8,10 +8,8 @@
 # Copyright (c) Meta Platforms, Inc. All Rights Reserved.
 
 
-from dataclasses import dataclass
-from typing import Callable, Optional, Tuple
+from typing import Callable, Optional
 
-import torch
 import torch.nn.functional as F
 from torch import nn