From a23488a65e3ab26c637b256284b12135f84c8803 Mon Sep 17 00:00:00 2001
From: TheVinz <37114928+TheVinz@users.noreply.github.com>
Date: Fri, 25 Jun 2021 16:40:25 +0200
Subject: [PATCH] Fix Trainer.restore_exp()

If self._model_optimizer is a dict, it has no load_state_dict() method. For loading each optimizer it is necessary to access the dictionary by using self._model.optimizer[key] first, and then use the load_state_dict() method.
---
 yapt/trainer/trainer.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/yapt/trainer/trainer.py b/yapt/trainer/trainer.py
index 332c16e..e2aa32e 100644
--- a/yapt/trainer/trainer.py
+++ b/yapt/trainer/trainer.py
@@ -355,7 +355,7 @@ def restore_exp(self):
 
         if is_dict(self._model.optimizer):
             for key in self._model.optimizer.keys():
-                self._model.optimizer.load_state_dict(
+                self._model.optimizer[key].load_state_dict(
                     checkpoint['optimizer_state_dict'][key])
         else:
             self._model.optimizer.load_state_dict(