Mark Transformer as deprecated

torchmdnet/models/torchmd_t.py

@@ -9,9 +9,10 @@
     act_class_mapping,
     scatter,
 )
+from torchmdnet.utils import deprecated_class
 
 
-
+@deprecated_class
 class TorchMD_T(nn.Module):
     r"""The TorchMD Transformer architecture.
 
@@ -155,7 +156,6 @@ def forward(
         s: Optional[Tensor] = None,
         q: Optional[Tensor] = None,
     ) -> Tuple[Tensor, Optional[Tensor], Tensor, Tensor, Tensor]:
-
         x = self.embedding(z)
 
         edge_index, edge_weight, _ = self.distance(pos, batch)

torchmdnet/utils.py

@@ -4,6 +4,8 @@
 import torch
 from os.path import dirname, join, exists
 from lightning_utilities.core.rank_zero import rank_zero_warn
+import functools
+import warnings
 
 # fmt: off
 # Atomic masses are based on:
@@ -217,3 +219,23 @@ def number(text):
 
 class MissingEnergyException(Exception):
     pass
+
+
+def deprecated_class(cls):
+    """Decorator to mark classes as deprecated."""
+    orig_init = cls.__init__
+
+    @functools.wraps(orig_init)
+    def wrapped_init(self, *args, **kwargs):
+        warnings.simplefilter(
+            "always", DeprecationWarning
+        )  # ensure all deprecation warnings are shown
+        warnings.warn(
+            f"{cls.__name__} is deprecated and will be removed in a future version.",
+            category=DeprecationWarning,
+            stacklevel=2,
+        )
+        orig_init(self, *args, **kwargs)
+
+    cls.__init__ = wrapped_init
+    return cls