Cast tensors in KVCache only when needed (#2017)

Author: Andrei-Aksionov

litgpt/model.py

@@ -791,8 +791,10 @@ def forward(self, input_pos: torch.Tensor, k: torch.Tensor, v: torch.Tensor) ->
 
         """
         # move the buffer to the activation dtype for when AMP is used
-        self.k = self.k.to(k.dtype)
-        self.v = self.v.to(v.dtype)
+        if self.k.dtype != k.dtype:
+            self.k = self.k.to(k.dtype)
+        if self.v.dtype != v.dtype:
+            self.v = self.v.to(v.dtype)
         # update the cache
         bs = k.size(0)
         k = batched_index_copy_(self.k[:bs, ...], -2, input_pos, k)