🏎 DeepSpeed Optimizer indexing

Zasder3 · Jun 12, 2021 · e18b7ca · e18b7ca
1 parent baebc38
commit e18b7ca
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/models/wrapper.py b/models/wrapper.py
@@ -66,6 +66,8 @@ def training_step(self, train_batch, idx):
             acc_t = (torch.argmax(image_logits, 0) == ground_truth).sum()
             self.log_dict({'loss': loss, 'acc': (acc_i + acc_t) / 2 / len(image)}, prog_bar=True)
 
+        if isinstance(optimizer, list):
+            optimizer = optimizer[0]
         optimizer.zero_grad()
 
         # image loss
@@ -207,6 +209,8 @@ def training_step(self, train_batch, idx):
             loss += (F.kl_div(image_logits_notemp * self.sink_temp, img_target) + F.kl_div(image_logits_notemp.t() * self.sink_temp, txt_target)) / 2 * self.kl_coeff
             self.log_dict({'loss': loss, 'acc': (acc_i + acc_t) / 2 / len(image)}, prog_bar=True)
 
+        if isinstance(optimizer, list):
+            optimizer = optimizer[0]
         optimizer.zero_grad()
 
         # image loss