move to() outside one level

FAIR-Chem · Aug 15, 2024 · 3bcc365 · 3bcc365
1 parent d0590af
commit 3bcc365
Showing 1 changed file with 3 additions and 4 deletions.
diff --git a/src/fairchem/core/trainers/base_trainer.py b/src/fairchem/core/trainers/base_trainer.py
@@ -605,11 +605,10 @@ def load_checkpoint(
                 target_key = key
 
             if target_key in self.normalizers:
-                mkeys = (
-                    self.normalizers[target_key]
-                    .load_state_dict(checkpoint["normalizers"][key])
-                    .to(map_location)
+                mkeys = self.normalizers[target_key].load_state_dict(
+                    checkpoint["normalizers"][key]
                 )
+                self.normalizers[target_key].to(map_location)
                 assert len(mkeys.missing_keys) == 0
                 assert len(mkeys.unexpected_keys) == 0