cleanup nits

Author: mirceamironenco

torchtune/modules/attention.py

@@ -270,7 +270,8 @@ def forward(
                 k = self.pos_embeddings(k, input_pos=input_pos)
 
             # k,v shape: [b, n_kv, s_y, h_d]
-            k, v = k.transpose(1, 2), v.transpose(1, 2)
+            k = k.transpose(1, 2)
+            v = v.transpose(1, 2)
 
             # Update key-value cache
             if self.kv_cache is not None and self.cache_enabled:

torchtune/modules/attention_utils.py

@@ -183,6 +183,7 @@ def _attention_call(
             dropout_p: float,
             is_causal: bool,
         ) -> torch.Tensor:
+
             # Flex attention uses the BlockMask
             # (https://github.com/pytorch/pytorch/blob/main/torch/nn/attention/flex_attention.py#L168)
             # instead of a traditional boolean tensor mask. If this is passed in,

torchtune/modules/kv_cache.py

@@ -51,7 +51,7 @@ def reset(self) -> None:
 
     @property
     def size(self) -> int:
-        return int(self.cache_pos[0].item())
+        return self.cache_pos[0].item()
 
     def update(
         self, k_val: torch.Tensor, v_val: torch.Tensor