From b851f472196269f00ac715b567bb0f7bf013123e Mon Sep 17 00:00:00 2001
From: Filippo Bigi <98903385+frostedoyster@users.noreply.github.com>
Date: Fri, 17 May 2024 11:34:28 +0200
Subject: [PATCH] Allow to access `wrapped_module` of exported models

---
 python/metatensor-torch/metatensor/torch/atomistic/model.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/python/metatensor-torch/metatensor/torch/atomistic/model.py b/python/metatensor-torch/metatensor/torch/atomistic/model.py
index d558b931f..af6d294c6 100644
--- a/python/metatensor-torch/metatensor/torch/atomistic/model.py
+++ b/python/metatensor-torch/metatensor/torch/atomistic/model.py
@@ -308,6 +308,7 @@ def __init__(
         else:
             raise ValueError(f"unknown dtype in capabilities: {capabilities.dtype}")
 
+    @torch.jit.export
     def wrapped_module(self) -> torch.nn.Module:
         """Get the module wrapped in this :py:class:`MetatensorAtomisticModel`"""
         return self._module