diff --git a/ai_models_panguweather/__init__.py b/ai_models_panguweather/__init__.py
index cae9874..9ab0fcc 100644
--- a/ai_models_panguweather/__init__.py
+++ b/ai_models_panguweather/__init__.py
@@ -5,4 +5,4 @@
 # granted to it by virtue of its status as an intergovernmental organisation
 # nor does it submit to any jurisdiction.
 
-__version__ = "0.0.3"
+__version__ = "0.0.4"
diff --git a/ai_models_panguweather/model.py b/ai_models_panguweather/model.py
index d5770b3..56c28c4 100644
--- a/ai_models_panguweather/model.py
+++ b/ai_models_panguweather/model.py
@@ -86,6 +86,8 @@ def run(self):
                 providers=self.providers,
             )
 
+        self.write_input_fields(fields_pl + fields_sfc)
+
         input_24, input_surface_24 = input, input_surface
 
         with self.stepper(6) as stepper:
diff --git a/setup.py b/setup.py
index e50c57e..74278a6 100644
--- a/setup.py
+++ b/setup.py
@@ -61,7 +61,7 @@ def has_gpu():
 
 setuptools.setup(
     name="ai-models-panguweather",
-    python_requires="<3.11",  # For now, does not support Python 3.11
+    # python_requires="<3.11",  # For now, does not support Python 3.11
     version=version,
     description="An ai-models plugin to run PanguWeather",
     long_description=read("README.md"),
@@ -74,7 +74,7 @@ def has_gpu():
     include_package_data=True,
     setup_requires=["GPUtil"],
     install_requires=[
-        "ai-models",
+        "ai-models>=0.3.6",
         "onnx",
         os.environ.get("ONNXRUNTIME", onnxruntime),
     ],