Fix OOM: batch sampling in micro-batches of 4

sample.py

@@ -168,8 +168,18 @@ def main():
     key = jax.random.PRNGKey(args.seed)
     print(f"Sampling {args.num_samples} configurations "
           f"(compiling on first call) …")
-    tokens = sample_batch(model, args.num_samples, key)
-    tokens = np.asarray(tokens)
+
+    # Sample in fixed-size micro-batches to avoid OOM on GPU.
+    # Changing _MICRO_BATCH triggers recompilation, so keep it constant.
+    _MICRO_BATCH = 4
+    all_tokens = []
+    n_calls = -(-args.num_samples // _MICRO_BATCH)  # ceiling division
+    for i in range(n_calls):
+        key, subkey = jax.random.split(key)
+        batch = sample_batch(model, _MICRO_BATCH, subkey)
+        all_tokens.append(np.asarray(batch))
+        print(f"  batch {i+1}/{n_calls} done")
+    tokens = np.concatenate(all_tokens)[:args.num_samples]
 
     grids = tokens_to_grids(tokens, lattice_size)   # (N, L, L), values {-1, +1}
     print(f"  shape: {grids.shape}  dtype: {grids.dtype}")