addressing torch.amp.GradScaler FutureWarnings

ml-struct-bio · Dec 17, 2024 · 25f1566 · 25f1566
1 parent 9989827
commit 25f1566
Show file tree

Hide file tree

Showing 4 changed files with 16 additions and 4 deletions.
diff --git a/cryodrgn/commands/abinit_het.py b/cryodrgn/commands/abinit_het.py
@@ -910,7 +910,10 @@ def main(args):
             model, optim = amp.initialize(model, optim, opt_level="O1")
         # mixed precision with pytorch (v1.6+)
         except:  # noqa: E722
-            scaler = torch.cuda.amp.grad_scaler.GradScaler()
+            try:
+                scaler = torch.amp.GradScaler("cuda")
+            except AttributeError:
+                scaler = torch.cuda.amp.grad_scaler.GradScaler()
 
     if args.load == "latest":
         args = get_latest(args)

diff --git a/cryodrgn/commands/abinit_homo.py b/cryodrgn/commands/abinit_homo.py
@@ -679,7 +679,10 @@ def main(args):
             model, optim = amp.initialize(model, optim, opt_level="O1")
         # mixed precision with pytorch (v1.6+)
         except:  # noqa: E722
-            scaler = torch.cuda.amp.grad_scaler.GradScaler()
+            try:
+                scaler = torch.amp.GradScaler("cuda")
+            except AttributeError:
+                scaler = torch.cuda.amp.grad_scaler.GradScaler()
 
     sorted_poses = []
     if args.load:

diff --git a/cryodrgn/commands/train_nn.py b/cryodrgn/commands/train_nn.py
@@ -513,7 +513,10 @@ def main(args: argparse.Namespace) -> None:
             model, optim = amp.initialize(model, optim, opt_level="O1")
         # Mixed precision with pytorch (v1.6+)
         except:  # noqa: E722
-            scaler = torch.cuda.amp.grad_scaler.GradScaler()
+            try:
+                scaler = torch.amp.GradScaler("cuda")
+            except AttributeError:
+                scaler = torch.cuda.amp.grad_scaler.GradScaler()
 
     # parallelize
     if args.multigpu and torch.cuda.device_count() > 1:

diff --git a/cryodrgn/commands/train_vae.py b/cryodrgn/commands/train_vae.py
@@ -864,7 +864,10 @@ def main(args: argparse.Namespace) -> None:
             model, optim = amp.initialize(model, optim, opt_level="O1")
         # mixed precision with pytorch (v1.6+)
         except:  # noqa: E722
-            scaler = torch.cuda.amp.grad_scaler.GradScaler()
+            try:
+                scaler = torch.amp.GradScaler("cuda")
+            except AttributeError:
+                scaler = torch.cuda.amp.grad_scaler.GradScaler()
 
     # restart from checkpoint
     if args.load: