[FIX] Run memory leak test separately

.github/workflows/test.yaml

@@ -1,14 +1,6 @@
 name: Test
 
-on:
-  push:
-    branches:
-      - '*'
-  pull_request:
-    branches:
-      - development
-      - master
-      - release
+on: [push, pull_request]
 
 jobs:
   tests:
@@ -34,10 +26,5 @@ jobs:
         make install-test
         pip install torch${{ matrix.pytorch-version }} torchvision
     - name: Run test
-      if: contains('refs/heads/master refs/heads/development refs/heads/release', github.ref)
       run: |
         make test
-    - name: Run test-light
-      if: contains('refs/heads/master refs/heads/development refs/heads/release', github.ref) != 1
-      run: |
-        make test-light

changelog.md

@@ -14,6 +14,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Fix CI (apply latest `black`, use latest `pytest` and `pytest-benchmark`)
   ([PR](https://github.com/f-dangel/backpack/pull/348))
+  - Follow-up fix, isolating a test that re-produces a memory leak.
+    Always run the full test suite
+    ([PR](https://github.com/f-dangel/backpack/pull/349))
 - Improve efficiency of Hessian-vector product
   ([PR](https://github.com/f-dangel/backpack/pull/341))
 

makefile

@@ -56,16 +56,20 @@ help:
 ###
 # Test coverage
 test:
-	@pytest -vx -rs --run-optional-tests=montecarlo --cov=backpack .
+	@pytest -vx -rs --run-optional-tests=montecarlo --cov=backpack . -k "test_memory_leak"
+	@pytest -vx -rs --run-optional-tests=montecarlo --cov=backpack . -k "not test_memory_leak"
 
 test-light:
-	@pytest -vx -rs --cov=backpack .
+	@pytest -vx -rs --cov=backpack . -k "test_memory_leak"
+	@pytest -vx -rs --cov=backpack . -k "not test_memory_leak"
 
 test-no-gpu:
-	@pytest -k "not cuda" -vx -rs --run-optional-tests=montecarlo --cov=backpack .
+	@pytest -k "(not cuda) and test_memory_leak" -vx -rs --run-optional-tests=montecarlo --cov=backpack .
+	@pytest -k "(not cuda) and (not test_memory_leak)" -vx -rs --run-optional-tests=montecarlo --cov=backpack .
 
 test-light-no-gpu:
-	@pytest -k "not cuda" -vx -rs --cov=backpack .
+	@pytest -k "(not cuda) and test_memory_leak" -vx -rs --cov=backpack .
+	@pytest -k "(not cuda) and (not test_memory_leak)" -vx -rs --cov=backpack .
 
 ###
 # Linter and autoformatter

test/test___init__.py

@@ -76,7 +76,7 @@ def memory_leak(device, context=None):
     if context is None:
         context = nullcontext
 
-    for _ in range(steps):
+    for s in range(steps):
         lossfunc = torch.nn.CrossEntropyLoss().to(device)
         lossfunc = extend(lossfunc)
 
@@ -87,6 +87,7 @@ def memory_leak(device, context=None):
             loss = lossfunc(model(X), y)  # this is what kills it
 
         memory = pytorch_current_memory_usage()
+        print(f"Step {s}, Memory delta: {(memory - memory_init) / 2**20:.2f} MiB.")
         if memory - memory_init > memory_leak_threshold:
             raise RuntimeError(
                 f"Memory leak detected: context={context}, device={device}"