diff --git a/distiller/apputils/checkpoint.py b/distiller/apputils/checkpoint.py
index d317bdead0e235680c9088152f8eb28770545b84..2757cd83e64e97eaba51c1067e1500edc4a6edaf 100755
--- a/distiller/apputils/checkpoint.py
+++ b/distiller/apputils/checkpoint.py
@@ -178,7 +178,7 @@ def load_checkpoint(model, chkpt_file, optimizer=None,
         try:
             if model.arch != checkpoint["arch"]:
                 raise ValueError("The model architecture does not match the checkpoint architecture")
-        except (NameError, KeyError):
+        except (AttributeError, KeyError):
             # One of the values is missing so we can't perform the comparison
             pass