fix probability (affects cpu results)

mehta-lab · Jan 14, 2025 · 6a88ec4 · 6a88ec4
1 parent dd00a36
commit 6a88ec4
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/viscy/scripts/bench_augmentations.py b/viscy/scripts/bench_augmentations.py
@@ -9,7 +9,7 @@
 seed_everything(42)
 
 # %%
-x = torch.rand(32, 2, 15, 512, 512, device="cpu")
+x = torch.rand(32, 2, 15, 512, 512, device="cuda")
 
 # %%
 monai_transform = RandAffine(
@@ -22,6 +22,7 @@
 kornia_transform = RandomAffine3D(
     degrees=(360.0, 0.0, 0.0),
     scale=((0.8, 1.2), (0.7, 1.3), (0.7, 1.3)),
+    p=1.0,
 )
 
 
@@ -50,15 +51,15 @@ def bench_kornia(x):
     globals=globals_injection,
     label="monai",
     setup="from __main__ import bench_monai",
-    num_threads=16
+    # num_threads=16,
 )
 
 kornia_timer = Timer(
     stmt="bench_kornia(x)",
     globals=globals_injection,
     label="kornia",
     setup="from __main__ import bench_kornia",
-    num_threads=16
+    # num_threads=16,
 )
 
 # %%