small naming changes

bashnick · Jan 23, 2023 · d8ce56f · d8ce56f
1 parent 7143710
commit d8ce56f
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/model.py b/model.py
@@ -4,9 +4,9 @@
 from utils import DEVICE
 
 
-class Head(nn.Module):
+class AttentionHead(nn.Module):
     """
-    One head of the self-attention
+    One head of the self-attention layer
     """
 
     def __init__(self, head_size, num_embed, block_size, dropout):
@@ -47,7 +47,7 @@ def __init__(self, num_heads, head_size, num_embed, block_size, dropout):
         super().__init__()
         self.heads = nn.ModuleList(
             [
-                Head(
+                AttentionHead(
                     head_size=head_size,
                     num_embed=num_embed,
                     block_size=block_size,