Merge 6ef9f77789 into 4a87a0d19f

Fix comment: rotary embeddings final dimension size

nanochat/gpt.py

@@ -244,7 +244,7 @@ class GPT(nn.Module):
     def forward(self, idx, targets=None, kv_cache=None, loss_reduction='mean'):
         B, T = idx.size()
 
-        # Grab the rotary embeddings for the current sequence length (they are of shape (1, seq_len, 1, head_dim))
+        # Grab the rotary embeddings for the current sequence length (they are of shape (1, seq_len, 1, head_dim/2))
         assert T <= self.cos.size(1), f"Sequence length grew beyond the rotary embeddings cache: {T} > {self.cos.size(1)}"
         assert idx.device == self.cos.device, f"Rotary embeddings and idx are on different devices: {idx.device} != {self.cos.device}"
         assert self.cos.dtype == torch.bfloat16, "Rotary embeddings must be in bfloat16"

tests/test_rustbpe.py

@@ -438,10 +438,12 @@ def enwik8_path():
     enwik8_local_path_zip = os.path.join(base_dir, "enwik8.zip")
     if not os.path.exists(enwik8_local_path):
         print(f"Downloading enwik8 to {enwik8_local_path_zip}")
-        import requests
-        response = requests.get(enwik8_url)
-        with open(enwik8_local_path_zip, "wb") as f:
-            f.write(response.content)
+        import urllib.request, urllib.error
+        try:
+            with urllib.request.urlopen(enwik8_url, timeout=30) as resp, open(enwik8_local_path_zip, "wb") as f:
+                f.write(resp.read())
+        except (urllib.error.URLError, urllib.error.HTTPError) as e:
+            pytest.skip(f"Network unavailable or download failed: {e}")
         with zipfile.ZipFile(enwik8_local_path_zip, "r") as zip_ref:
             zip_ref.extractall(base_dir)
         print(f"Unzipped enwik8 to {enwik8_local_path}")