more changes

nod-ai · Feb 25, 2025 · 9a02c48 · 9a02c48
1 parent d664d88
commit 9a02c48
Show file tree

Hide file tree

Showing 3 changed files with 2 additions and 2 deletions.
diff --git a/sharktank/sharktank/layers/paged_llama_attention_block.py b/sharktank/sharktank/layers/paged_llama_attention_block.py
@@ -103,7 +103,6 @@ def forward(
     ):
         assert bool(start_index is not None) ^ bool(embedding_batch_mask is not None)
 
-        print("c", h[0, :, 0])
         x = self.attn_norm(h)
         bs, batch_seq_len, _ = x.shape
 

diff --git a/sharktank/sharktank/layers/token_embedding.py b/sharktank/sharktank/layers/token_embedding.py
@@ -24,4 +24,5 @@ def __init__(
         self.dtype = dtype
 
     def forward(self, input: torch.Tensor):
+        print("embedding", self.weight[0, 0])
         return ops.embedding_lookup(input, self.weight, dtype=self.dtype)
diff --git a/sharktank/sharktank/models/grok/toy_grok.py b/sharktank/sharktank/models/grok/toy_grok.py
@@ -19,7 +19,6 @@
 
 
 def generate(seed):
-    torch.manual_seed(seed)
     dtype = torch.float16
     block_seq_stride = 16
     max_blocks = 8
@@ -52,6 +51,7 @@ def generate(seed):
         attention_dtype=dtype,
     )
 
+    torch.manual_seed(seed)
     theta = make_random_grok_theta(
         config=config,
         vocab_size=vocabulary_size,