Merge pull request #632 from mv1388/fix-grad-clip-multi-optimizer-cal…

…lback Fix grad clip multi optimizer callback
mv1388 · Oct 31, 2020 · 8d0ee4f · 8d0ee4f
2 parents 00d98d9 + 7bf2126
commit 8d0ee4f
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 2 deletions.
diff --git a/aitoolbox/torchtrain/callbacks/gradient.py b/aitoolbox/torchtrain/callbacks/gradient.py
@@ -3,6 +3,7 @@
 import torch
 
 from aitoolbox.torchtrain.callbacks.abstract import AbstractCallback, AbstractExperimentCallback
+from aitoolbox.torchtrain.multi_loss_optim import MultiOptimizer
 from aitoolbox.experiment.local_save.local_results_save import BaseLocalResultsSaver
 from aitoolbox.experiment.result_reporting.report_generator import GradientPlotter
 from aitoolbox.cloud.AWS.results_save import BaseResultsSaver as BaseResultsS3Saver
@@ -41,7 +42,11 @@ def __init__(self, max_grad_value):
 
     def on_after_gradient_update(self, optimizer_idx):
         if self.train_loop_obj.use_amp:
-            self.train_loop_obj.amp_scaler.unscale_(self.train_loop_obj.optimizer)
+            optimizer = self.train_loop_obj.optimizer
+            if isinstance(optimizer, MultiOptimizer):
+                optimizer = optimizer[optimizer_idx]
+
+            self.train_loop_obj.amp_scaler.unscale_(optimizer)
 
         torch.nn.utils.clip_grad_value_(self.train_loop_obj.model.parameters(), self.max_grad_value)
 
@@ -60,7 +65,11 @@ def __init__(self, max_grad_norm, **kwargs):
 
     def on_after_gradient_update(self, optimizer_idx):
         if self.train_loop_obj.use_amp:
-            self.train_loop_obj.amp_scaler.unscale_(self.train_loop_obj.optimizer)
+            optimizer = self.train_loop_obj.optimizer
+            if isinstance(optimizer, MultiOptimizer):
+                optimizer = optimizer[optimizer_idx]
+
+            self.train_loop_obj.amp_scaler.unscale_(optimizer)
 
         torch.nn.utils.clip_grad_norm_(self.train_loop_obj.model.parameters(), self.max_grad_norm, **self.kwargs)
 

diff --git a/dist/aitoolbox-1.2.0-py3-none-any.whl b/dist/aitoolbox-1.2.0-py3-none-any.whl
diff --git a/dist/aitoolbox-1.2.0.tar.gz b/dist/aitoolbox-1.2.0.tar.gz