diff --git a/TrainingExtensions/torch/src/python/setup.py b/TrainingExtensions/torch/src/python/setup.py
index e39504faf34..e0c6e3a8692 100644
--- a/TrainingExtensions/torch/src/python/setup.py
+++ b/TrainingExtensions/torch/src/python/setup.py
@@ -103,7 +103,7 @@ def has_ext_modules(self):
     license="NOTICE.txt",
     long_description=(PACKAGING_DIR / "README.txt").read_text(),
     name="AimetTorch",
-    package_data={"": ["*.json", "*.html"], },
+    package_data={"": ["*.json", "*.html", "*.js", "*.xml"], },
     package_dir={"": ".", },
     packages=find_namespace_packages(where=".", exclude=["build"]),
     platforms="x86",