fix deprecated custom forward method

mehta-lab · Aug 30, 2024 · 4babd35 · 4babd35
1 parent 9b6e3f2
commit 4babd35
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/viscy/light/engine.py b/viscy/light/engine.py
@@ -103,7 +103,7 @@ def __init__(
         self.l2_alpha = l2_alpha
         self.ms_dssim_alpha = ms_dssim_alpha
 
-    @torch.cuda.amp.custom_fwd(cast_inputs=torch.float32)
+    @torch.amp.custom_fwd(cast_inputs=torch.float32)
     def forward(self, preds, target):
         loss = 0
         if self.l1_alpha: