Fix torchscript issue with additive models and PET

metatensor · Dec 16, 2024 · 4d27b1e · 4d27b1e
1 parent 9e7e5f9
commit 4d27b1e
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/src/metatrain/experimental/pet/model.py b/src/metatrain/experimental/pet/model.py
@@ -186,9 +186,9 @@ def forward(
             # at evaluation, we also add the additive contributions
             for additive_model in self.additive_models:
                 outputs_for_additive_model: Dict[str, ModelOutput] = {}
-                for name, output in outputs.items():
-                    if name in additive_model.outputs:
-                        outputs_for_additive_model[name] = output
+                for output_name, output_options in outputs.items():
+                    if output_name in additive_model.outputs:
+                        outputs_for_additive_model[output_name] = output_options
                 additive_contributions = additive_model(
                     systems,
                     outputs_for_additive_model,