Using fused step for optimizers

ancestor-mithril · Apr 29, 2024 · a0716d3 · a0716d3
1 parent 9d2fcbe
commit a0716d3
Showing 1 changed file with 13 additions and 1 deletion.
diff --git a/utils/optimizer.py b/utils/optimizer.py
@@ -1,3 +1,4 @@
+import torch
 from torch import optim
 from torch.optim import Optimizer
 
@@ -7,4 +8,15 @@ def init_optimizer(args, parameters) -> Optimizer:
     momentum = 0.9 if not hasattr(args, 'momentum') else args.momentum
     weight_decay = 5e-4 if not hasattr(args, 'weight_decay') else args.weight_decay
 
-    return optim.SGD(parameters, lr=args.lr, momentum=momentum, weight_decay=weight_decay)
+    kwargs = {
+        "lr": args.lr,
+        "momentum": momentum,
+        "weight_decay": weight_decay,
+    }
+    if torch.torch_version.TorchVersion(torch.__version__) >= '2.3.0':
+        kwargs['fused'] = True
+    else:
+        import warnings
+        warnings.warn("Upgrade torch to support fused optimizers")
+
+    return optim.SGD(parameters, **kwargs)