remove uneeded prints

54e8cad2d6e63779ad0c0955406a5565e389ace6 · 2024-07-14 23:31:18 -0700 · Alex Cheema

Files touched

M exo/inference/mlx/models/sharded_llama.py

Diff

commit 54e8cad2d6e63779ad0c0955406a5565e389ace6
Author: Alex Cheema <alexcheema123@gmail.com>
Date:   Sun Jul 14 23:31:18 2024 -0700

    remove uneeded prints
---
 exo/inference/mlx/models/sharded_llama.py | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

diff --git a/exo/inference/mlx/models/sharded_llama.py b/exo/inference/mlx/models/sharded_llama.py
index c2d9deda..267d177e 100644
--- a/exo/inference/mlx/models/sharded_llama.py
+++ b/exo/inference/mlx/models/sharded_llama.py
@@ -90,8 +90,6 @@ class Attention(nn.Module):
     ) -> mx.array:
         B, L, D = x.shape
 
-        print("q_proj: ", self.q_proj)
-        print("x: ", x.shape)
         queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x)
 
         # Prepare the queries, keys and values for the attention computation
@@ -192,8 +190,7 @@ class LlamaModel(nn.Module):
         if cache is None:
             cache = [None] * len(self.layers)
 
-        for i, (layer, c) in enumerate(zip(self.layers, cache)):
-            print(f"layer: {i}")
+        for layer, c in zip(self.layers, cache):
             h = layer(h, mask, cache=c)
 
         if self.args.shard.is_last_layer():

← c6912055 empty space · back to Exo · add llama-3-70b to the examples b6595bac →