import in getters, so registry is populated

Author: baberabb

lm_eval/api/registry.py

@@ -486,7 +486,7 @@ def decorate(cls):
     return decorate
 
 
-def get_model(model_name):
+def get_model(model_name: str):
     """Get a model class by name.
 
     Args:

lm_eval/evaluator.py

@@ -8,7 +8,6 @@
 from typing import TYPE_CHECKING, List, Optional, Union
 
 import numpy as np
-import torch
 
 import lm_eval.api.metrics
 import lm_eval.api.registry
@@ -33,6 +32,7 @@
     hash_dict_images,
     hash_string,
     positional_deprecated,
+    set_torch_seed,
     setup_logging,
     simple_parse_args_string,
     wrap_text,
@@ -193,7 +193,7 @@ def simple_evaluate(
 
     if torch_random_seed is not None:
         seed_message.append(f"Setting torch manual seed to {torch_random_seed}")
-        torch.manual_seed(torch_random_seed)
+        set_torch_seed(torch_random_seed)
 
     if fewshot_random_seed is not None:
         seed_message.append(f"Setting fewshot manual seed to {fewshot_random_seed}")
@@ -553,6 +553,8 @@ def evaluate(
             requests[reqtype].append(instance)
 
         if lm.world_size > 1:
+            import torch
+
             instances_rnk = torch.tensor(len(task._instances), device=lm.device)
             gathered_item = (
                 lm.accelerator.gather(instances_rnk).cpu().detach().numpy().tolist()
@@ -661,6 +663,8 @@ def evaluate(
                     task_output.sample_metrics[(metric, filter_key)].append(value)
 
     if WORLD_SIZE > 1:
+        import torch
+
         # if multigpu, then gather data across all ranks to rank 0
         # first gather logged samples across all ranks
         for task_output in eval_tasks:

lm_eval/utils.py

@@ -840,3 +840,9 @@ def _request_with_retries(method, url, **kwargs):
         return False
 
     return True
+
+
+def set_torch_seed(seed: int):
+    import torch
+
+    torch.manual_seed(seed)