pytest

Author: HojeChun

nff/tests/conftest.py

@@ -2,10 +2,19 @@
 
 import pytest
 import torch
+import builtins
 
 torch.set_num_threads(int(os.getenv("OMP_NUM_THREADS", "1")))
 
 
+def pytest_sessionstart(session):
+    """
+    Pytorch 2.6 torch.load compatibility
+    """
+    add_safe = getattr(torch.serialization, "add_safe_globals", None)
+    if callable(add_safe):
+        add_safe([builtins.slice])
+
 def pytest_addoption(parser):
     parser.addoption("--device", action="store", default="cpu", help="Whether to use the CPU or GPU for the tests")