update

ecmwf · May 28, 2024 · be95203 · be95203
1 parent 86c9444
commit be95203
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 1 deletion.
diff --git a/pyproject.toml b/pyproject.toml
@@ -44,6 +44,7 @@ dependencies = [
     "anemoi-utils>=0.2.0",
     "semantic-version",
     "pyyaml",
+    "torch",
 ]
 
 [project.optional-dependencies]

diff --git a/src/anemoi/inference/__main__.py b/src/anemoi/inference/__main__.py
@@ -20,7 +20,7 @@
 LOG = logging.getLogger(__name__)
 
 
-def main():
+def create_parser():
     parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
 
     parser.add_argument(
@@ -41,6 +41,11 @@ def main():
         command_parser = subparsers.add_parser(name, help=command.__doc__)
         command.add_arguments(command_parser)
 
+    return parser
+
+
+def main():
+    parser = create_parser()
     args = parser.parse_args()
 
     if args.version:

diff --git a/src/anemoi/inference/runner.py b/src/anemoi/inference/runner.py
@@ -68,6 +68,7 @@ def run(
         progress_callback=ignore,
         add_ensemble_dimension=False,
     ):
+        import torch
 
         if autocast is None:
             autocast = self.checkpoint.precision