diff --git a/src/cryo_sbi/wpa_simulator/cryo_em_simulator.py b/src/cryo_sbi/wpa_simulator/cryo_em_simulator.py
index b097a1e..4ba5525 100644
--- a/src/cryo_sbi/wpa_simulator/cryo_em_simulator.py
+++ b/src/cryo_sbi/wpa_simulator/cryo_em_simulator.py
@@ -12,7 +12,7 @@
 from cryo_sbi.utils.micrograph_utils import RandomMicrographPatches
 from cryo_sbi.utils.image_utils import NormalizeIndividual
 
-
+#@torch.no_grad()
 def cryo_em_simulator(
     models,
     index,