Add version hyperparameter to the checkpoints.

docs/source/installation.rst

@@ -54,7 +54,7 @@ It is recommended to install the same version as the one used by torch.
 
 .. code-block:: shell
 
-   conda install -c conda-forge cuda-nvcc cuda-libraries-dev cuda-version "gxx<12" pytorch=*=*cuda*
+   conda install -c conda-forge cuda-nvcc cuda-libraries-dev cuda-version "gxx<12" "pytorch=*=*cuda*"
 
 
 .. warning:: gxx<12 is required due to a `bug in GCC+CUDA12 <https://github.com/pybind/pybind11/issues/4606>`_ that prevents pybind11 from compiling correctly

environment.yml

@@ -17,3 +17,4 @@ dependencies:
   - pytest
   - psutil
   - gxx<12
+  - versioneer==0.28

setup.cfg

@@ -0,0 +1,7 @@
+[versioneer]
+VCS = git
+style = pep440
+versionfile_source = torchmdnet/_version.py
+versionfile_build = torchmdnet/_version.py
+tag_prefix = 
+parentdir_prefix = torchmdnet-

setup.py

@@ -2,26 +2,24 @@
 # Distributed under the MIT License.
 # (See accompanying file README.md file or copy at http://opensource.org/licenses/MIT)
 
-import subprocess
 from setuptools import setup, find_packages
 import torch
-from torch.utils.cpp_extension import BuildExtension, CUDAExtension, include_paths, CppExtension
+from torch.utils.cpp_extension import (
+    BuildExtension,
+    CUDAExtension,
+    include_paths,
+    CppExtension,
+)
+import versioneer
 import os
 
-try:
-    version = (
-        subprocess.check_output(["git", "describe", "--abbrev=0", "--tags"])
-        .strip()
-        .decode("utf-8")
-    )
-except:
-    print("Failed to retrieve the current version, defaulting to 0")
-    version = "0"
 # If CPU_ONLY is defined
 force_cpu_only = os.environ.get("CPU_ONLY", None) is not None
 use_cuda = torch.cuda._is_compiled() if not force_cpu_only else False
+
+
 def set_torch_cuda_arch_list():
-    """ Set the CUDA arch list according to the architectures the current torch installation was compiled for.
+    """Set the CUDA arch list according to the architectures the current torch installation was compiled for.
     This function is a no-op if the environment variable TORCH_CUDA_ARCH_LIST is already set or if torch was not compiled with CUDA support.
     """
     if not os.environ.get("TORCH_CUDA_ARCH_LIST"):
@@ -32,31 +30,36 @@ def set_torch_cuda_arch_list():
             formatted_versions += "+PTX"
             os.environ["TORCH_CUDA_ARCH_LIST"] = formatted_versions
 
+
 set_torch_cuda_arch_list()
 
-extension_root= os.path.join("torchmdnet", "extensions")
-neighbor_sources=["neighbors_cpu.cpp"]
+extension_root = os.path.join("torchmdnet", "extensions")
+neighbor_sources = ["neighbors_cpu.cpp"]
 if use_cuda:
     neighbor_sources.append("neighbors_cuda.cu")
-neighbor_sources = [os.path.join(extension_root, "neighbors", source) for source in neighbor_sources]
+neighbor_sources = [
+    os.path.join(extension_root, "neighbors", source) for source in neighbor_sources
+]
 
 ExtensionType = CppExtension if not use_cuda else CUDAExtension
 extensions = ExtensionType(
-    name='torchmdnet.extensions.torchmdnet_extensions',
+    name="torchmdnet.extensions.torchmdnet_extensions",
     sources=[os.path.join(extension_root, "extensions.cpp")] + neighbor_sources,
     include_dirs=include_paths(),
-    define_macros=[('WITH_CUDA', 1)] if use_cuda else [],
+    define_macros=[("WITH_CUDA", 1)] if use_cuda else [],
 )
 
 if __name__ == "__main__":
+    buildext = BuildExtension.with_options(no_python_abi_suffix=True, use_ninja=False)
     setup(
         name="torchmd-net",
-        version=version,
         packages=find_packages(),
-        ext_modules=[extensions],
-        cmdclass={
-            'build_ext': BuildExtension.with_options(no_python_abi_suffix=True, use_ninja=False)},
         include_package_data=True,
-        entry_points={"console_scripts": ["torchmd-train = torchmdnet.scripts.train:main"]},
+        entry_points={
+            "console_scripts": ["torchmd-train = torchmdnet.scripts.train:main"]
+        },
         package_data={"torchmdnet": ["extensions/torchmdnet_extensions.so"]},
+        ext_modules=[extensions],
+        version=versioneer.get_version(),
+        cmdclass=versioneer.get_cmdclass({"build_ext": buildext}),
     )

torchmdnet/__init__.py

@@ -0,0 +1,3 @@
+from torchmdnet import _version
+
+__version__ = _version.get_versions()["version"]