enable any dict and namespace in hparams (#1847)

Lightning-AI · May 15, 2020 · b84b024 · b84b024
1 parent e95e1d7
commit b84b024
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/pytorch_lightning/trainer/training_io.py b/pytorch_lightning/trainer/training_io.py
@@ -344,10 +344,11 @@ def dump_checkpoint(self):
 
         if hasattr(model, "hparams") and model.hparams is not None:
             parsing.clean_namespace(model.hparams)
-            checkpoint['hparams_type'] = model.hparams.__class__.__name__
-            if checkpoint['hparams_type'] == 'dict':
+            if isinstance(model.hparams, dict):
+                checkpoint['hparams_type'] = 'dict'
                 checkpoint['hparams'] = model.hparams
-            elif checkpoint['hparams_type'] == 'Namespace':
+            elif isinstance(model.hparams, Namespace):
+                checkpoint['hparams_type'] = 'Namespace'
                 checkpoint['hparams'] = vars(model.hparams)
             else:
                 raise ValueError(