evo2_fix.patch

diff --git a/evo2/test/test_evo2_generation.py b/evo2/test/test_evo2_generation.py
index 789ecb3..2345c75 100644
--- a/evo2/test/test_evo2_generation.py
+++ b/evo2/test/test_evo2_generation.py
@@ -4,6 +4,7 @@ from importlib import resources
 from pathlib import Path
 from typing import List, Optional, Union
 import numpy as np
+import time
 
 import torch
 
@@ -65,6 +66,8 @@ def generate_and_score(*, sequences, model, generations_per_prompt=5, n_tokens=5
         target = targets[i]
 
         with torch.inference_mode():
+            if torch.cuda.is_available(): torch.cuda.synchronize()
+            elapsed_time = -time.perf_counter()
             generated = model.generate(
                 prompt_seqs=[prompt],
                 n_tokens=n_tokens,
@@ -72,7 +75,10 @@ def generate_and_score(*, sequences, model, generations_per_prompt=5, n_tokens=5
                 top_k=top_k,
                 top_p=top_p,
             )
-            
+            if torch.cuda.is_available(): torch.cuda.synchronize()
+            elapsed_time += time.perf_counter()
+            print(f"[{i}] Time for model.generate: {elapsed_time:.3f} s")
+
             decoded_seq = generated.sequences[0]  # Assuming generate returns list of sequences
             score = calculate_sequence_identity(decoded_seq, target)
             scores.append(score)
@@ -94,6 +100,7 @@ def main():
     parser = argparse.ArgumentParser(description="Test Evo2 Model Generation")
     parser.add_argument("--model_name", choices=['evo2_7b', 'evo2_40b', 'evo2_1b_base'], default='evo2_7b',
                        help="Model to test (supports evo2_7b, evo2_40b, evo2_1b_base)")
+    parser.add_argument("--local_path", type=str, default=None)
     
     args = parser.parse_args()
     
@@ -101,8 +108,8 @@ def main():
     torch.manual_seed(1)
     torch.cuda.manual_seed(1)
         
-    model = Evo2(args.model_name)
-    
+    model = Evo2(args.model_name, local_path=args.local_path)
+
     # Test parameters: greedy sampling of 500 tokens
     test_params = {
         'n_tokens': 500,
@@ -140,4 +147,4 @@ def main():
         print(f"\nTest Failed: Expected {expected_score}%, got {mean_score}%")
 
 if __name__ == "__main__":
-    main()
\ No newline at end of file
+    main()
diff --git a/pyproject.toml b/pyproject.toml
index fb1c1e6..6c7bd64 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -3,7 +3,7 @@ name = "evo2"
 version = "0.3.0"
 description = "Genome modeling across all domains of life"
 readme = "README.md"
-requires-python = ">=3.11,<3.13"
+requires-python = ">=3.10,<3.13"
 license = "Apache-2.0"
 authors = [
     {name = "Evo 2 Team"},