amp.load_state_dict fix

2020-05-03 04:54:10 -07:00 · 2020-05-03 04:54:10 -07:00 · 807011a821
--- a/archai/common/apex_utils.py
+++ b/archai/common/apex_utils.py
@ -238,6 +238,6 @@ class ApexUtils:

    def load_state_dict(self, state_dict):
        if self.is_mixed():
-            self._amp.load_state_dict()
+            self._amp.load_state_dict(state_dict)