Removed the leftover GPU NL warning (#107)

Author: frostedoyster

src/upet/__init__.py

@@ -3,6 +3,7 @@
 import warnings
 
 import torch
+from warp._src import utils as wp_utils
 
 from ._models import get_upet, save_upet
 
@@ -14,10 +15,22 @@
     message="The .grad attribute of a Tensor that is not a leaf Tensor",
 )
 
+# we want to suppress a further warning from nvalchemi's usage of warp
+# warp uses an internal warn() helper: we wrap it.
+_orig_warn = wp_utils.warn
+
+
+def _warn_filtered(message, category=None, stacklevel=1):
+    if category is DeprecationWarning and "warp.vec" in str(message):
+        return
+    return _orig_warn(message, category=category, stacklevel=stacklevel)
+
+
+wp_utils.warn = _warn_filtered
+
 # Disable static fusion. Besides the fact that atomistic batches have variable
 # sizes, statically fused CUDA kernels cannot allocate new tensors at runtime,
 # causing "Global alloc not supported yet" errors (cuda 13+) at the time of writing
 torch.jit.set_fusion_strategy([("DYNAMIC", 10)])
 
-
 __all__ = ["get_upet", "save_upet"]